azure: move to default fsspec loop (#6323)

dvc/fs/azure.py

@@ -1,10 +1,8 @@
-import asyncio
 import logging
 import os
-import sys
 import threading
-from contextlib import contextmanager
 
+from fsspec.asyn import fsspec_loop
 from fsspec.utils import infer_storage_options
 from funcy import cached_property, memoize, wrap_prop
 
@@ -23,47 +21,6 @@
 )
 
 
-@contextmanager
-def _temp_event_loop():
-    """When trying to initialize azure filesystem instances
-    with DefaultCredentials, the authentication process requires
-    to have an access to a separate event loop. The normal calls
-    are run in a separate loop managed by the fsspec, but the
-    DefaultCredentials assumes that the callee is managing their
-    own event loop. This function checks whether is there any
-    event loop set, and if not it creates a temporary event loop
-    and set it. After the context is finalized, it restores the
-    original event loop back (if there is any)."""
-
-    try:
-        original_loop = asyncio.get_event_loop()
-        original_policy = asyncio.get_event_loop_policy()
-    except RuntimeError:
-        original_loop = None
-        original_policy = None
-
-    # From 3.8>= and onwards, asyncio changed the default
-    # loop policy for windows to use proactor loops instead
-    # of selector based ones. Due to that, proxied connections
-    # doesn't work with the aiohttp and this is most likely an
-    # upstream bug that needs to be solved outside of DVC. Until
-    # such issue is resolved, we need to manage this;
-    # https://github.com/aio-libs/aiohttp/issues/4536
-    if sys.version_info >= (3, 8) and os.name == "nt":
-        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
-
-    loop = original_loop or asyncio.new_event_loop()
-
-    try:
-        asyncio.set_event_loop(loop)
-        yield
-    finally:
-        if original_loop is None:
-            loop.close()
-        asyncio.set_event_loop(original_loop)
-        asyncio.set_event_loop_policy(original_policy)
-
-
 class AzureAuthError(DvcException):
     pass
 
@@ -149,9 +106,10 @@ def _prepare_credentials(self, **config):
             and not any_secondary
             and not config.get("allow_anonymous_login", False)
         ):
-            login_info["credential"] = DefaultAzureCredential(
-                exclude_interactive_browser_credential=False
-            )
+            with fsspec_loop():
+                login_info["credential"] = DefaultAzureCredential(
+                    exclude_interactive_browser_credential=False
+                )
 
         for login_method, required_keys in [  # noqa
             ("connection string", ["connection_string"]),
@@ -182,19 +140,10 @@ def fs(self):
         from azure.core.exceptions import AzureError
 
         try:
-            with _temp_event_loop():
-                file_system = AzureBlobFileSystem(**self.fs_args)
+            return AzureBlobFileSystem(**self.fs_args)
         except (ValueError, AzureError) as e:
             raise AzureAuthError(
                 f"Authentication to Azure Blob Storage via {self.login_method}"
                 " failed.\nLearn more about configuration settings at"
                 f" {format_link('https://man.dvc.org/remote/modify')}"
             ) from e
-
-        return file_system
-
-    def open(
-        self, path_info, mode="r", **kwargs
-    ):  # pylint: disable=arguments-differ
-        with _temp_event_loop():
-            return self.fs.open(self._with_bucket(path_info), mode=mode)

tests/unit/fs/test_azure.py

@@ -1,9 +1,6 @@
-import asyncio
-from concurrent.futures import ThreadPoolExecutor
-
 import pytest
 
-from dvc.fs.azure import AzureAuthError, AzureFileSystem, _temp_event_loop
+from dvc.fs.azure import AzureAuthError, AzureFileSystem
 from dvc.path_info import PathInfo
 
 container_name = "container-name"
@@ -50,31 +47,6 @@ def test_info(tmp_dir, azure):
     assert hash_.strip("'").strip('"') == hash_
 
 
-def test_temp_event_loop():
-    def procedure():
-        loop = asyncio.get_event_loop()
-        loop.run_until_complete(asyncio.sleep(0))
-        return "yeey"
-
-    def wrapped_procedure():
-        with _temp_event_loop():
-            return procedure()
-
-        # it should clean the loop after
-        # exitting the context.
-        with pytest.raises(RuntimeError):
-            asyncio.get_event_loop()
-
-    with ThreadPoolExecutor(max_workers=1) as executor:
-        future = executor.submit(procedure)
-
-        with pytest.raises(RuntimeError):
-            future.result()
-
-        future = executor.submit(wrapped_procedure)
-        assert future.result() == "yeey"
-
-
 def test_azure_login_methods():
     def get_login_method(config):
         fs = AzureFileSystem(**config)