Make downloader import inside the function

microsoft · Dec 6, 2023 · 76edae1 · 76edae1
1 parent c7f4e4f
commit 76edae1
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/src/promptflow/promptflow/azure/operations/_run_operations.py b/src/promptflow/promptflow/azure/operations/_run_operations.py
@@ -64,7 +64,6 @@
 from promptflow.azure._restclient.flow.models import SetupFlowSessionAction
 from promptflow.azure._restclient.flow_service_caller import FlowServiceCaller
 from promptflow.azure._utils.gerneral import get_user_alias_from_credential
-from promptflow.azure.operations._async_run_downloader import AsyncRunDownloader
 from promptflow.azure.operations._flow_operations import FlowOperations
 from promptflow.exceptions import UserErrorException
 
@@ -1020,6 +1019,8 @@ def download(
         """
         import platform
 
+        from promptflow.azure.operations._async_run_downloader import AsyncRunDownloader
+
         run = Run._validate_and_return_run_name(run)
         if output is None:
             # default to be "~/.promptflow/.runs" folder
@@ -1045,7 +1046,7 @@ def download(
             # On Windows seems to be a problem with EventLoopPolicy, use this snippet to work around it
             asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
 
-        async_run_allowing_running_loop(asyncio.run(run_downloader.download()))
+        async_run_allowing_running_loop(run_downloader.download)
         result_path = run_folder.resolve().as_posix()
         logger.info(f"Successfully downloaded run {run!r} to {result_path!r}.")
         return result_path
diff --git a/src/promptflow/tests/sdk_cli_azure_test/e2etests/test_run_operations.py b/src/promptflow/tests/sdk_cli_azure_test/e2etests/test_run_operations.py
@@ -21,7 +21,6 @@
 from promptflow._utils.flow_utils import get_flow_lineage_id
 from promptflow.azure import PFClient
 from promptflow.azure._entities._flow import Flow
-from promptflow.azure.operations._async_run_downloader import AsyncRunDownloader
 from promptflow.exceptions import UserErrorException
 
 from .._azure_utils import DEFAULT_TEST_TIMEOUT, PYTEST_TIMEOUT_METHOD
@@ -799,6 +798,8 @@ def test_vnext_workspace_base_url(self):
 
     @pytest.mark.skipif(condition=not is_live(), reason="need to fix recording")
     def test_download_run(self, pf):
+        from promptflow.azure.operations._async_run_downloader import AsyncRunDownloader
+
         run = "c619f648-c809-4545-9f94-f67b0a680706"
 
         expected_files = [