Add async run logs and status polling

 * Add AsyncRunClient support for get_logs, inspect, wait_for_condition, and watch_statuses.
 * Share status wait condition helpers while keeping sync and async polling control flow explicit.

Author: md-mq

cli/polyaxon/_client/run.py

@@ -999,23 +999,25 @@ def get_statuses(
         except (ApiException, HTTPError) as e:
             raise PolyaxonClientException("Api error: %s" % e) from e
 
-    def _wait_for_condition(self, statuses: Optional[List[str]] = None):
-        statuses = to_list(statuses, check_none=True)
+    def _should_stop_waiting(self, last_status, statuses) -> bool:
+        if statuses:
+            return last_status in statuses
+        return LifeCycle.is_done(last_status)
 
-        def condition():
-            if statuses:
-                return last_status in statuses
-            return LifeCycle.is_done(last_status)
+    def _is_retryable_status_error(self, error: ApiException) -> bool:
+        return error.status in {500, 502, 503, 504}
 
+    def _wait_for_condition(self, statuses: Optional[List[str]] = None):
+        statuses = to_list(statuses, check_none=True)
         last_status = None
-        while not condition():
+        while not self._should_stop_waiting(last_status, statuses):
             if last_status:
                 time.sleep(settings.CLIENT_CONFIG.watch_interval)
             try:
                 last_status, _conditions = self.get_statuses(last_status)
                 yield last_status, _conditions
             except ApiException as e:
-                if e.status in {500, 502, 503, 504}:
+                if self._is_retryable_status_error(e):
                     yield last_status, []
                 else:
                     raise e
@@ -3435,25 +3437,82 @@ async def get_statuses(
         except ApiException as e:
             raise PolyaxonClientException("Api error: %s" % e) from e
 
-    @async_client_handler(check_no_op=True, check_offline=True)
-    async def wait_for_condition(self, *args, **kwargs):
-        self._raise_sync_only("wait_for_condition")
+    async def _wait_for_condition(self, statuses: Optional[List[str]] = None):
+        statuses = to_list(statuses, check_none=True)
+        last_status = None
+        while not self._should_stop_waiting(last_status, statuses):
+            if last_status:
+                await asyncio.sleep(settings.CLIENT_CONFIG.watch_interval)
+            try:
+                last_status, conditions = await self.get_statuses(last_status)
+                yield last_status, conditions
+            except ApiException as e:
+                if self._is_retryable_status_error(e):
+                    yield last_status, []
+                else:
+                    raise e
 
     @async_client_handler(check_no_op=True, check_offline=True)
-    async def watch_statuses(self, *args, **kwargs):
-        self._raise_sync_only("watch_statuses")
+    async def wait_for_condition(
+        self,
+        statuses: Optional[List[str]] = None,
+        print_status: bool = False,
+        live_update: Any = None,
+    ):
+        async for status, _conditions in self._wait_for_condition(statuses):
+            self._run_data.status = status  # type: ignore
+            if print_status:
+                print("Last received status: {}\n".format(status))
+            if live_update:
+                latest_status = Printer.add_status_color(
+                    {"status": status}, status_key="status"
+                )
+                live_update.update(status="{}\n".format(latest_status["status"]))
+
+    async def watch_statuses(self, statuses: Optional[List[str]] = None):
+        if self._no_op or self._is_offline:
+            return
+        async for status, conditions in self._wait_for_condition(statuses):
+            self._run_data.status = status  # type: ignore
+            yield status, conditions
 
     @async_client_handler(check_no_op=True, check_offline=True)
-    async def get_logs(self, *args, **kwargs):
-        self._raise_sync_only("get_logs")
+    async def get_logs(self, last_file=None, last_time=None) -> "V1Logs":
+        if not self.settings:
+            await self.refresh_data()
+        await self._use_agent_host()
+        params = get_logs_params(
+            last_file=last_file,
+            last_time=last_time,
+            connection=self.artifacts_store,
+            kind=self.run_data.kind,
+        )
+        return await self.client.runs_v1.get_run_logs(
+            self.namespace,
+            self.owner,
+            self.project,
+            self.run_uuid,
+            **params,
+        )
 
     @async_client_handler(check_no_op=True, check_offline=True)
     async def watch_logs(self, *args, **kwargs):
         self._raise_sync_only("watch_logs")
 
     @async_client_handler(check_no_op=True, check_offline=True)
-    async def inspect(self, *args, **kwargs):
-        self._raise_sync_only("inspect")
+    async def inspect(self):
+        if not self.settings:
+            await self.refresh_data()
+        await self._use_agent_host()
+        params = get_streams_params(connection=self.artifacts_store, status=self.status)
+        return await self.client.runs_v1.inspect_run(
+            self.namespace,
+            self.owner,
+            self.project,
+            self.run_uuid,
+            self.run_data.kind,
+            **params,
+        )
 
     @async_client_handler(check_no_op=True, check_offline=True)
     async def shell(self, *args, **kwargs):

cli/tests/test_client/test_async_run_client.py

@@ -2,6 +2,7 @@
 from mock import mock
 import pytest
 
+from polyaxon import settings
 from polyaxon._client.run import AsyncRunClient, RunClient
 from polyaxon._schemas.lifecycle import (
     V1ProjectVersionKind,
@@ -162,7 +163,10 @@ def test_async_run_client_rejects_sync_client():
 def test_async_run_client_method_surface_is_async():
     for method in ASYNC_METHODS:
         assert method in AsyncRunClient.__dict__
-        assert inspect.iscoroutinefunction(getattr(AsyncRunClient, method))
+        if method == "watch_statuses":
+            assert inspect.isasyncgenfunction(getattr(AsyncRunClient, method))
+        else:
+            assert inspect.iscoroutinefunction(getattr(AsyncRunClient, method))
 
 
 def test_async_run_client_local_helpers_stay_shared():
@@ -329,6 +333,109 @@ async def test_status_methods_await_api_without_async_req():
     assert "async_req" not in sdk_client.runs_v1.create_run_status.call_args[1]
 
 
+@pytest.mark.asyncio
+async def test_wait_for_condition_updates_status():
+    patch_settings()
+    sdk_client = AsyncPolyaxonClientMock()
+    sdk_client.runs_v1.get_run_statuses = AsyncMock(
+        return_value=mock.Mock(
+            status=V1Statuses.RUNNING,
+            status_conditions=[
+                V1StatusCondition.model_construct(type=V1Statuses.RUNNING)
+            ],
+        )
+    )
+    client = make_client(sdk_client)
+
+    await client.wait_for_condition(statuses=[V1Statuses.RUNNING])
+
+    assert client.status == V1Statuses.RUNNING
+    sdk_client.runs_v1.get_run_statuses.assert_called_once_with(
+        OWNER,
+        PROJECT,
+        RUN_UUID,
+    )
+
+
+@pytest.mark.asyncio
+async def test_watch_statuses_yields_until_done_status():
+    patch_settings()
+    settings.CLIENT_CONFIG.watch_interval = 0
+    sdk_client = AsyncPolyaxonClientMock()
+    running_condition = V1StatusCondition.model_construct(type=V1Statuses.RUNNING)
+    succeeded_condition = V1StatusCondition.model_construct(type=V1Statuses.SUCCEEDED)
+    sdk_client.runs_v1.get_run_statuses = AsyncMock(
+        side_effect=[
+            mock.Mock(
+                status=V1Statuses.RUNNING,
+                status_conditions=[running_condition],
+            ),
+            mock.Mock(
+                status=V1Statuses.SUCCEEDED,
+                status_conditions=[running_condition, succeeded_condition],
+            ),
+        ]
+    )
+    client = make_client(sdk_client)
+
+    statuses = []
+    async for status, _conditions in client.watch_statuses():
+        statuses.append(status)
+
+    assert statuses == [V1Statuses.RUNNING, V1Statuses.SUCCEEDED]
+    assert client.status == V1Statuses.SUCCEEDED
+    assert sdk_client.runs_v1.get_run_statuses.call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_get_logs_awaits_refresh_and_api():
+    patch_settings()
+    sdk_client = AsyncPolyaxonClientMock()
+    response = mock.Mock()
+    sdk_client.runs_v1.get_run = AsyncMock(return_value=make_run())
+    sdk_client.runs_v1.get_run_logs = AsyncMock(return_value=response)
+    client = make_client(sdk_client)
+    client._run_data.settings = None
+
+    result = await client.get_logs(last_file="last.log", last_time="123")
+
+    assert result is response
+    sdk_client.runs_v1.get_run.assert_called_once_with(OWNER, PROJECT, RUN_UUID)
+    assert sdk_client.runs_v1.get_run_logs.call_args[0] == (
+        "test-namespace",
+        OWNER,
+        PROJECT,
+        RUN_UUID,
+    )
+    assert "async_req" not in sdk_client.runs_v1.get_run_logs.call_args[1]
+
+
+@pytest.mark.asyncio
+async def test_inspect_awaits_refresh_and_api():
+    patch_settings()
+    sdk_client = AsyncPolyaxonClientMock()
+    response = {"pods": {}}
+    sdk_client.runs_v1.get_run = AsyncMock(
+        return_value=make_run(status=V1Statuses.RUNNING)
+    )
+    sdk_client.runs_v1.inspect_run = AsyncMock(return_value=response)
+    client = make_client(sdk_client)
+    client._run_data.settings = None
+
+    result = await client.inspect()
+
+    assert result is response
+    sdk_client.runs_v1.get_run.assert_called_once_with(OWNER, PROJECT, RUN_UUID)
+    assert sdk_client.runs_v1.inspect_run.call_args[0] == (
+        "test-namespace",
+        OWNER,
+        PROJECT,
+        RUN_UUID,
+        None,
+    )
+    assert "async_req" not in sdk_client.runs_v1.inspect_run.call_args[1]
+
+
 @pytest.mark.asyncio
 async def test_events_metrics_and_lineage_methods_await_api():
     patch_settings()
@@ -504,7 +611,6 @@ async def test_promote_methods_use_async_project_client_with_injected_client():
 @pytest.mark.parametrize(
     "method,args",
     [
-        ("get_logs", ()),
         ("upload_artifact", ("file.txt",)),
         ("download_artifacts", ()),
         ("persist_run", ("/tmp/run",)),