From 5dda6fd9749a886585cebf9afc288ebc46f00429 Mon Sep 17 00:00:00 2001 From: Daniel Hollas Date: Tue, 12 Sep 2023 14:57:40 +0100 Subject: [PATCH] Refactor: Delay import of heavy packages to speed up import time (#6116) The importing of packages `urllib`, `yaml` and `pgsu` are moved from top-level to inside the scopes where they are needed. This significantly improves the load time of the `aiida` package and its subpackages. The import of `Transport` in the `aiida.engine` package also slows down imports, but it is only used for type checking, so its import is placed inside the `if TYPE_CHECKING` guard. Finally, the `DEFAULT_DBINFO`, `Postgres` and `PostgresConnectionMode` objects of the `aiida.manage.external.postgres` package are no longer exposed on the top-level as this also slows down imports. This is a breaking change technically, but these resources should not be used by downstream packages. --- aiida/cmdline/commands/cmd_archive.py | 3 ++- aiida/cmdline/commands/cmd_code.py | 3 ++- aiida/cmdline/commands/cmd_devel.py | 14 ++++++++++++-- aiida/cmdline/commands/cmd_rabbitmq.py | 3 ++- aiida/cmdline/params/options/config.py | 3 ++- aiida/cmdline/params/options/main.py | 2 +- aiida/cmdline/params/types/path.py | 8 ++++++-- aiida/cmdline/utils/echo.py | 5 ++++- aiida/engine/daemon/execmanager.py | 6 ++++-- aiida/engine/processes/builder.py | 4 ++-- aiida/engine/processes/calcjobs/monitors.py | 4 +++- aiida/engine/transports.py | 8 +++++--- aiida/manage/__init__.py | 3 --- aiida/manage/external/__init__.py | 4 ---- aiida/manage/external/postgres.py | 2 -- aiida/manage/external/rmq/client.py | 3 ++- aiida/manage/external/rmq/utils.py | 3 ++- aiida/schedulers/scheduler.py | 4 +++- docs/source/nitpick-exceptions | 1 + 19 files changed, 53 insertions(+), 30 deletions(-) diff --git a/aiida/cmdline/commands/cmd_archive.py b/aiida/cmdline/commands/cmd_archive.py index 70f9e7b02b..1c933c0e64 100644 --- a/aiida/cmdline/commands/cmd_archive.py +++ b/aiida/cmdline/commands/cmd_archive.py @@ -14,7 +14,6 @@ from pathlib import Path import traceback from typing import List, Tuple -import urllib.request import click from click_spinner import spinner @@ -432,6 +431,8 @@ def _import_archive_and_migrate( :param try_migration: whether to try a migration if the import raises `IncompatibleStorageSchema` """ + import urllib.request + from aiida.common.folders import SandboxFolder from aiida.tools.archive.abstract import get_format from aiida.tools.archive.imports import import_archive as _import_archive diff --git a/aiida/cmdline/commands/cmd_code.py b/aiida/cmdline/commands/cmd_code.py index eee6b2901c..305b1078e3 100644 --- a/aiida/cmdline/commands/cmd_code.py +++ b/aiida/cmdline/commands/cmd_code.py @@ -13,7 +13,6 @@ import click import tabulate -import yaml from aiida.cmdline.commands.cmd_verdi import verdi from aiida.cmdline.groups.dynamic import DynamicEntryPointCommandGroup @@ -242,6 +241,8 @@ def show(code): @with_dbenv() def export(code, output_file): """Export code to a yaml file.""" + import yaml + code_data = {} for key in code.get_cli_options().keys(): diff --git a/aiida/cmdline/commands/cmd_devel.py b/aiida/cmdline/commands/cmd_devel.py index ccd4be4cff..d123030f15 100644 --- a/aiida/cmdline/commands/cmd_devel.py +++ b/aiida/cmdline/commands/cmd_devel.py @@ -66,8 +66,18 @@ def devel_check_undesired_imports(): loaded_modules = 0 for modulename in [ - 'asyncio', 'requests', 'plumpy', 'disk_objectstore', 'paramiko', 'seekpath', 'CifFile', 'ase', 'pymatgen', - 'spglib', 'pymysql' + 'asyncio', + 'requests', + 'plumpy', + 'disk_objectstore', + 'paramiko', + 'seekpath', + 'CifFile', + 'ase', + 'pymatgen', + 'spglib', + 'pymysql', + 'yaml', ]: if modulename in sys.modules: echo.echo_warning(f'Detected loaded module "{modulename}"') diff --git a/aiida/cmdline/commands/cmd_rabbitmq.py b/aiida/cmdline/commands/cmd_rabbitmq.py index 000176928a..16c9d95dde 100644 --- a/aiida/cmdline/commands/cmd_rabbitmq.py +++ b/aiida/cmdline/commands/cmd_rabbitmq.py @@ -18,7 +18,6 @@ import click import tabulate import wrapt -import yaml from aiida.cmdline.commands.cmd_devel import verdi_devel from aiida.cmdline.params import arguments, options @@ -146,6 +145,8 @@ def with_manager(wrapped, _, args, kwargs): @with_manager def cmd_server_properties(manager): """List the server properties.""" + import yaml + data = {} for key, value in manager.get_communicator().server_properties.items(): data[key] = value.decode('utf-8') if isinstance(value, bytes) else value diff --git a/aiida/cmdline/params/options/config.py b/aiida/cmdline/params/options/config.py index c0694e9f4a..d8db0465b4 100644 --- a/aiida/cmdline/params/options/config.py +++ b/aiida/cmdline/params/options/config.py @@ -24,7 +24,6 @@ import typing as t import click -import yaml from .overridable import OverridableOption @@ -33,6 +32,8 @@ def yaml_config_file_provider(handle, cmd_name): # pylint: disable=unused-argument """Read yaml config file from file handle.""" + import yaml + return yaml.safe_load(handle) diff --git a/aiida/cmdline/params/options/main.py b/aiida/cmdline/params/options/main.py index f7afd3ea50..062c365cc7 100644 --- a/aiida/cmdline/params/options/main.py +++ b/aiida/cmdline/params/options/main.py @@ -9,9 +9,9 @@ ########################################################################### """Module with pre-defined reusable commandline options that can be used as `click` decorators.""" import click -from pgsu import DEFAULT_DSN as DEFAULT_DBINFO # pylint: disable=no-name-in-module from aiida.common.log import LOG_LEVELS, configure_logging +from aiida.manage.external.postgres import DEFAULT_DBINFO from aiida.manage.external.rmq import BROKER_DEFAULTS from .. import types diff --git a/aiida/cmdline/params/types/path.py b/aiida/cmdline/params/types/path.py index de016e42e9..33de95513c 100644 --- a/aiida/cmdline/params/types/path.py +++ b/aiida/cmdline/params/types/path.py @@ -10,8 +10,6 @@ """Click parameter types for paths.""" import os from socket import timeout -import urllib.error -import urllib.request import click @@ -88,6 +86,9 @@ def convert(self, value, param, ctx): def checks_url(self, url, param, ctx): """Check whether URL is reachable within timeout.""" + import urllib.error + import urllib.request + try: with urllib.request.urlopen(url, timeout=self.timeout_seconds): pass @@ -123,6 +124,9 @@ def convert(self, value, param, ctx): def get_url(self, url, param, ctx): """Retrieve file from URL.""" + import urllib.error + import urllib.request + try: return urllib.request.urlopen(url, timeout=self.timeout_seconds) # pylint: disable=consider-using-with except (urllib.error.URLError, urllib.error.HTTPError, timeout): diff --git a/aiida/cmdline/utils/echo.py b/aiida/cmdline/utils/echo.py index 6f49babee5..394c5143d6 100644 --- a/aiida/cmdline/utils/echo.py +++ b/aiida/cmdline/utils/echo.py @@ -16,7 +16,6 @@ from typing import Any, Optional import click -import yaml CMDLINE_LOGGER = logging.getLogger('verdi') @@ -225,11 +224,15 @@ def default_jsondump(data): def _format_yaml(dictionary, sort_keys=True): """Return a dictionary formatted as a string using the YAML format.""" + import yaml + return yaml.dump(dictionary, sort_keys=sort_keys) def _format_yaml_expanded(dictionary, sort_keys=True): """Return a dictionary formatted as a string using the expanded YAML format.""" + import yaml + return yaml.dump(dictionary, sort_keys=sort_keys, default_flow_style=False) diff --git a/aiida/engine/daemon/execmanager.py b/aiida/engine/daemon/execmanager.py index 3fa5612bba..5ff8a432a8 100644 --- a/aiida/engine/daemon/execmanager.py +++ b/aiida/engine/daemon/execmanager.py @@ -21,9 +21,9 @@ import pathlib import shutil from tempfile import NamedTemporaryFile -from typing import Any, List from typing import Mapping as MappingType from typing import Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, List from aiida.common import AIIDA_LOGGER, exceptions from aiida.common.datastructures import CalcInfo @@ -35,7 +35,9 @@ from aiida.orm.utils.log import get_dblogger_extra from aiida.repository.common import FileType from aiida.schedulers.datastructures import JobState -from aiida.transports import Transport + +if TYPE_CHECKING: + from aiida.transports import Transport REMOTE_WORK_DIRECTORY_LOST_FOUND = 'lost+found' diff --git a/aiida/engine/processes/builder.py b/aiida/engine/processes/builder.py index 8a50b1426f..9362288d26 100644 --- a/aiida/engine/processes/builder.py +++ b/aiida/engine/processes/builder.py @@ -13,8 +13,6 @@ from typing import TYPE_CHECKING, Any, Type from uuid import uuid4 -import yaml - from aiida.engine.processes.ports import PortNamespace from aiida.orm import Dict, Node from aiida.orm.nodes.data.base import BaseType @@ -245,6 +243,8 @@ def process_class(self) -> Type['Process']: def _repr_pretty_(self, p, _) -> str: # pylint: disable=invalid-name """Pretty representation for in the IPython console and notebooks.""" + import yaml + return p.text( f'Process class: {self._process_class.__name__}\n' f'Inputs:\n{yaml.safe_dump(json.JSONDecoder().decode(PrettyEncoder().encode(self)))}' diff --git a/aiida/engine/processes/calcjobs/monitors.py b/aiida/engine/processes/calcjobs/monitors.py index 1866539e95..9bc08b6a0f 100644 --- a/aiida/engine/processes/calcjobs/monitors.py +++ b/aiida/engine/processes/calcjobs/monitors.py @@ -13,7 +13,9 @@ from aiida.common.log import AIIDA_LOGGER from aiida.orm import CalcJobNode, Dict from aiida.plugins import BaseFactory -from aiida.transports import Transport + +if t.TYPE_CHECKING: + from aiida.transports import Transport LOGGER = AIIDA_LOGGER.getChild(__name__) diff --git a/aiida/engine/transports.py b/aiida/engine/transports.py index 1a046aea65..d0eae4ed14 100644 --- a/aiida/engine/transports.py +++ b/aiida/engine/transports.py @@ -13,10 +13,12 @@ import contextvars import logging import traceback -from typing import Awaitable, Dict, Hashable, Iterator, Optional +from typing import TYPE_CHECKING, Awaitable, Dict, Hashable, Iterator, Optional from aiida.orm import AuthInfo -from aiida.transports import Transport + +if TYPE_CHECKING: + from aiida.transports import Transport _LOGGER = logging.getLogger(__name__) @@ -55,7 +57,7 @@ def loop(self) -> asyncio.AbstractEventLoop: return self._loop @contextlib.contextmanager - def request_transport(self, authinfo: AuthInfo) -> Iterator[Awaitable[Transport]]: + def request_transport(self, authinfo: AuthInfo) -> Iterator[Awaitable['Transport']]: """ Request a transport from an authinfo. Because the client is not allowed to request a transport immediately they will instead be given back a future diff --git a/aiida/manage/__init__.py b/aiida/manage/__init__.py index a958f61b6a..857c19cbc6 100644 --- a/aiida/manage/__init__.py +++ b/aiida/manage/__init__.py @@ -35,13 +35,10 @@ 'CURRENT_CONFIG_VERSION', 'Config', 'ConfigValidationError', - 'DEFAULT_DBINFO', 'MIGRATIONS', 'ManagementApiConnectionError', 'OLDEST_COMPATIBLE_CONFIG_VERSION', 'Option', - 'Postgres', - 'PostgresConnectionMode', 'Profile', 'RabbitmqManagementClient', 'check_and_migrate_config', diff --git a/aiida/manage/external/__init__.py b/aiida/manage/external/__init__.py index c773109351..7ec4dcbd01 100644 --- a/aiida/manage/external/__init__.py +++ b/aiida/manage/external/__init__.py @@ -14,15 +14,11 @@ # yapf: disable # pylint: disable=wildcard-import -from .postgres import * from .rmq import * __all__ = ( 'BROKER_DEFAULTS', - 'DEFAULT_DBINFO', 'ManagementApiConnectionError', - 'Postgres', - 'PostgresConnectionMode', 'RabbitmqManagementClient', 'get_launch_queue_name', 'get_message_exchange_name', diff --git a/aiida/manage/external/postgres.py b/aiida/manage/external/postgres.py index 3a31d2248f..cd2b471667 100644 --- a/aiida/manage/external/postgres.py +++ b/aiida/manage/external/postgres.py @@ -23,8 +23,6 @@ if TYPE_CHECKING: from aiida.manage.configuration import Profile -__all__ = ('Postgres', 'PostgresConnectionMode', 'DEFAULT_DBINFO') - # The last placeholder is for adding privileges of the user _CREATE_USER_COMMAND = 'CREATE USER "{}" WITH PASSWORD \'{}\' {}' _DROP_USER_COMMAND = 'DROP USER "{}"' diff --git a/aiida/manage/external/rmq/client.py b/aiida/manage/external/rmq/client.py index e56fb5945d..476c63111a 100644 --- a/aiida/manage/external/rmq/client.py +++ b/aiida/manage/external/rmq/client.py @@ -3,7 +3,6 @@ from __future__ import annotations import typing as t -from urllib.parse import quote from aiida.common.exceptions import AiidaException @@ -49,6 +48,8 @@ def format_url(self, url: str, url_params: dict[str, str] | None = None) -> str: automatically inserted and should not be specified. :returns: The complete URL. """ + from urllib.parse import quote + url_params = url_params or {} url_params['virtual_host'] = self._virtual_host if self._virtual_host else '/' url_params = {key: quote(value, safe='') for key, value in url_params.items()} diff --git a/aiida/manage/external/rmq/utils.py b/aiida/manage/external/rmq/utils.py index 314f3946d3..6357250cb6 100644 --- a/aiida/manage/external/rmq/utils.py +++ b/aiida/manage/external/rmq/utils.py @@ -1,6 +1,5 @@ # -*- coding: utf-8 -*- """Utilites for RabbitMQ.""" -from urllib.parse import urlencode, urlunparse from . import defaults @@ -25,6 +24,8 @@ def get_rmq_url(protocol=None, username=None, password=None, host=None, port=Non :param kwargs: remaining keyword arguments that will be encoded as query parameters. :returns: the connection URL string. """ + from urllib.parse import urlencode, urlunparse + if 'heartbeat' not in kwargs: kwargs['heartbeat'] = defaults.BROKER_DEFAULTS.heartbeat diff --git a/aiida/schedulers/scheduler.py b/aiida/schedulers/scheduler.py index 6ea15af960..d98e343d67 100644 --- a/aiida/schedulers/scheduler.py +++ b/aiida/schedulers/scheduler.py @@ -19,7 +19,9 @@ from aiida.common.lang import classproperty from aiida.engine.processes.exit_code import ExitCode from aiida.schedulers.datastructures import JobInfo, JobResource, JobTemplate, JobTemplateCodeInfo -from aiida.transports import Transport + +if t.TYPE_CHECKING: + from aiida.transports import Transport __all__ = ('Scheduler', 'SchedulerError', 'SchedulerParsingError') diff --git a/docs/source/nitpick-exceptions b/docs/source/nitpick-exceptions index 60e0c89b89..d9286bc889 100644 --- a/docs/source/nitpick-exceptions +++ b/docs/source/nitpick-exceptions @@ -8,6 +8,7 @@ py:class callable py:class function py:class traceback py:class NoneType +py:class MappingType py:class AbstractContextManager py:class BinaryIO py:class IO