Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

azure: move to default fsspec loop #6323

Merged
merged 1 commit into from
Jul 18, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
63 changes: 6 additions & 57 deletions dvc/fs/azure.py
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -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

Expand Down Expand Up @@ -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"]),
Expand Down Expand Up @@ -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)
30 changes: 1 addition & 29 deletions tests/unit/fs/test_azure.py
Original file line number Diff line number Diff line change
@@ -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"
Expand Down Expand Up @@ -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)
Expand Down