Add async run artifact downloads

Author: md-mq

cli/polyaxon/_client/run.py

@@ -3687,16 +3687,159 @@ async def get_artifact(
         )
 
     @async_client_handler(check_no_op=True, check_offline=True)
-    async def download_artifact_for_lineage(self, *args, **kwargs):
-        self._raise_sync_only("download_artifact_for_lineage")
+    async def download_artifact_for_lineage(
+        self,
+        lineage: V1RunArtifact,
+        force: bool = False,
+        path_to: Optional[str] = None,
+    ):
+        if not self.run_uuid:
+            return
+
+        if not self.settings:
+            await self.refresh_data()
+        await self._use_agent_host()
+
+        lineage_path = lineage.path or ""
+        summary = lineage.summary or {}
+        is_event = summary.get("is_event")
+        has_step = summary.get("step")
+
+        if self.run_uuid in lineage_path:
+            lineage_path = os.path.relpath(lineage_path, self.run_uuid)
+
+        if V1ArtifactKind.is_single_file_event(lineage.kind):
+            return await self.download_artifact(
+                path=lineage_path,
+                force=force,
+                path_to=path_to,
+            )
+
+        if V1ArtifactKind.is_single_or_multi_file_event(lineage.kind):
+            if is_event or has_step:
+                if not self.settings:
+                    await self.refresh_data()
+                url = get_proxy_run_url(
+                    service=STREAMS_V1_LOCATION,
+                    namespace=self.namespace,
+                    owner=self.owner,
+                    project=self.project,
+                    run_uuid=self.run_uuid,
+                    subpath="events/{}".format(lineage.kind),
+                )
+                url = absolute_uri(url=url, host=self.client.config.host)
+                params = get_streams_params(
+                    connection=self.artifacts_store,
+                    force=force,
+                )
+                params.update({"names": lineage.name, "pkg_assets": True})
+
+                # TODO: Update with AsyncPolyaxonStore is done
+                return await asyncio.to_thread(
+                    self.store.download_file,
+                    url=url,
+                    path=self.run_uuid,
+                    use_filepath=False,
+                    extract_path=path_to,
+                    path_to=path_to,
+                    params=params,
+                    untar=True,
+                )
+            if V1ArtifactKind.is_file_or_dir(lineage.kind):
+                return await self.download_artifacts(
+                    path=lineage_path,
+                    path_to=path_to,
+                    check_path=True,
+                )
+            return await self.download_artifact(
+                path=lineage_path,
+                force=force,
+                path_to=path_to,
+            )
+
+        if V1ArtifactKind.is_file(lineage.kind):
+            return await self.download_artifact(
+                path=lineage_path,
+                force=force,
+                path_to=path_to,
+            )
+
+        if V1ArtifactKind.is_dir(lineage.kind):
+            return await self.download_artifacts(path=lineage_path, path_to=path_to)
+
+        if V1ArtifactKind.is_file_or_dir(lineage.kind):
+            return await self.download_artifacts(
+                path=lineage_path,
+                path_to=path_to,
+                check_path=True,
+            )
 
     @async_client_handler(check_no_op=True, check_offline=True)
-    async def download_artifact(self, *args, **kwargs):
-        self._raise_sync_only("download_artifact")
+    async def download_artifact(
+        self,
+        path: str,
+        force: bool = False,
+        path_to: Optional[str] = None,
+    ):
+        if not self.settings:
+            await self.refresh_data()
+        await self._use_agent_host()
+
+        url = get_proxy_run_url(
+            service=STREAMS_V1_LOCATION,
+            namespace=self.namespace,
+            owner=self.owner,
+            project=self.project,
+            run_uuid=self.run_uuid,
+            subpath="artifact",
+        )
+        url = absolute_uri(url=url, host=self.client.config.host)
+        params = get_streams_params(connection=self.artifacts_store, force=force)
+        return await asyncio.to_thread(
+            self.store.download_file,
+            url=url,
+            path=path,
+            path_to=path_to,
+            params=params,
+        )
 
     @async_client_handler(check_no_op=True, check_offline=True)
-    async def download_artifacts(self, *args, **kwargs):
-        self._raise_sync_only("download_artifacts")
+    async def download_artifacts(
+        self,
+        path: str = "",
+        path_to: Optional[str] = None,
+        untar: bool = True,
+        delete_tar: bool = True,
+        extract_path: Optional[str] = None,
+        check_path: bool = False,
+    ):
+        if not self.settings:
+            await self.refresh_data()
+        await self._use_agent_host()
+
+        url = get_proxy_run_url(
+            service=STREAMS_V1_LOCATION,
+            namespace=self.namespace,
+            owner=self.owner,
+            project=self.project,
+            run_uuid=self.run_uuid,
+            subpath="artifacts",
+        )
+        url = absolute_uri(url=url, host=self.client.config.host)
+        params = get_streams_params(connection=self.artifacts_store)
+        if check_path:
+            params["check_path"] = True
+
+        return await asyncio.to_thread(
+            self.store.download_file,
+            url=url,
+            path=path,
+            untar=untar,
+            path_to=path_to,
+            delete_tar=delete_tar and untar,
+            extract_path=extract_path,
+            params=params,
+        )
 
     @async_client_handler(check_no_op=True, check_offline=True)
     async def upload_artifact(self, *args, **kwargs):

cli/tests/test_client/test_async_run_client.py

@@ -575,6 +575,163 @@ async def test_get_artifact_awaits_refresh_when_settings_missing():
     )
 
 
+@pytest.mark.asyncio
+async def test_download_artifact_methods_call_store_download_file():
+    patch_settings()
+    sdk_client = AsyncPolyaxonClientMock()
+    sdk_client.config = mock.Mock(host="http://api", no_op=None, is_offline=None)
+    client = make_client(sdk_client)
+    client._run_data = make_run()
+    client._store = mock.Mock()
+    client._store.download_file = mock.Mock(side_effect=["/tmp/file", "/tmp/archive"])
+
+    file_path = await client.download_artifact(
+        "outputs/model.pkl",
+        force=True,
+        path_to="/tmp/downloads",
+    )
+    archive_path = await client.download_artifacts(
+        "outputs",
+        path_to="/tmp/downloads",
+        untar=False,
+        extract_path="/tmp/extract",
+        check_path=True,
+    )
+
+    assert file_path == "/tmp/file"
+    assert archive_path == "/tmp/archive"
+    file_kwargs = client._store.download_file.call_args_list[0][1]
+    archive_kwargs = client._store.download_file.call_args_list[1][1]
+    assert file_kwargs["path"] == "outputs/model.pkl"
+    assert file_kwargs["path_to"] == "/tmp/downloads"
+    assert file_kwargs["params"] == {"force": True}
+    assert file_kwargs["url"].endswith(
+        "/streams/v1/test-namespace/test-owner/test-project/runs/{}/artifact".format(
+            RUN_UUID
+        )
+    )
+    assert archive_kwargs["path"] == "outputs"
+    assert archive_kwargs["path_to"] == "/tmp/downloads"
+    assert archive_kwargs["untar"] is False
+    assert archive_kwargs["delete_tar"] is False
+    assert archive_kwargs["extract_path"] == "/tmp/extract"
+    assert archive_kwargs["params"] == {"check_path": True}
+    assert archive_kwargs["url"].endswith(
+        "/streams/v1/test-namespace/test-owner/test-project/runs/{}/artifacts".format(
+            RUN_UUID
+        )
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "lineage,expected_method,expected_kwargs",
+    [
+        (
+            V1RunArtifact.model_construct(
+                kind=V1ArtifactKind.METRIC,
+                path=f"{RUN_UUID}/events/metric/loss",
+            ),
+            "download_artifact",
+            {
+                "path": "events/metric/loss",
+                "force": True,
+                "path_to": "/tmp/downloads",
+            },
+        ),
+        (
+            V1RunArtifact.model_construct(
+                kind=V1ArtifactKind.MODEL,
+                path="models/model.pkl",
+            ),
+            "download_artifacts",
+            {
+                "path": "models/model.pkl",
+                "path_to": "/tmp/downloads",
+                "check_path": True,
+            },
+        ),
+        (
+            V1RunArtifact.model_construct(
+                kind=V1ArtifactKind.DIR,
+                path="outputs",
+            ),
+            "download_artifacts",
+            {
+                "path": "outputs",
+                "path_to": "/tmp/downloads",
+            },
+        ),
+    ],
+)
+async def test_download_artifact_for_lineage_routes_to_download_helpers(
+    lineage,
+    expected_method,
+    expected_kwargs,
+):
+    patch_settings()
+    sdk_client = AsyncPolyaxonClientMock()
+    client = make_client(sdk_client)
+    client._run_data = make_run()
+    client.download_artifact = AsyncMock(return_value="file")
+    client.download_artifacts = AsyncMock(return_value="artifacts")
+
+    result = await client.download_artifact_for_lineage(
+        lineage,
+        force=True,
+        path_to="/tmp/downloads",
+    )
+
+    if expected_method == "download_artifact":
+        assert result == "file"
+        client.download_artifact.assert_called_once_with(**expected_kwargs)
+        assert client.download_artifacts.call_count == 0
+    else:
+        assert result == "artifacts"
+        client.download_artifacts.assert_called_once_with(**expected_kwargs)
+        assert client.download_artifact.call_count == 0
+
+
+@pytest.mark.asyncio
+async def test_download_artifact_for_lineage_downloads_event_package():
+    patch_settings()
+    sdk_client = AsyncPolyaxonClientMock()
+    sdk_client.config = mock.Mock(host="http://api", no_op=None, is_offline=None)
+    client = make_client(sdk_client)
+    client._run_data = make_run()
+    client._store = mock.Mock()
+    client._store.download_file = mock.Mock(return_value="/tmp/events")
+    lineage = V1RunArtifact.model_construct(
+        kind=V1ArtifactKind.MODEL,
+        name="model",
+        path="events/model",
+        summary={"is_event": True},
+    )
+
+    result = await client.download_artifact_for_lineage(
+        lineage,
+        force=True,
+        path_to="/tmp/downloads",
+    )
+
+    assert result == "/tmp/events"
+    kwargs = client._store.download_file.call_args[1]
+    assert kwargs["path"] == RUN_UUID
+    assert kwargs["use_filepath"] is False
+    assert kwargs["extract_path"] == "/tmp/downloads"
+    assert kwargs["path_to"] == "/tmp/downloads"
+    assert kwargs["untar"] is True
+    assert kwargs["params"] == {
+        "force": True,
+        "names": "model",
+        "pkg_assets": True,
+    }
+    assert kwargs["url"].endswith(
+        "/streams/v1/test-namespace/test-owner/test-project/runs/"
+        "{}/events/model".format(RUN_UUID)
+    )
+
+
 @pytest.mark.asyncio
 async def test_run_action_methods_await_api_without_async_req():
     patch_settings()
@@ -840,9 +997,6 @@ async def test_promote_methods_use_async_project_client_with_injected_client():
         ("upload_artifact", ("file.txt",)),
         ("upload_artifacts_dir", ("outputs",)),
         ("upload_artifacts", (["file.txt"],)),
-        ("download_artifact_for_lineage", (V1RunArtifact.model_construct(),)),
-        ("download_artifact", ("file.txt",)),
-        ("download_artifacts", ()),
         ("persist_run", ("/tmp/run",)),
         ("push_offline_run", ("/tmp/run",)),
         ("get_runs_as_hiplot", ()),