Skip to content

Commit

Permalink
Refactor: Delay import of heavy packages to speed up import time (#6106)
Browse files Browse the repository at this point in the history
The importing of packages `disk_objectstore`, `jsonschema`, `requests`,
`plumpy` and `paramiko` 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 `ProcessLauncher` utility had to be removed from resources that are
exposed at a higher package level because it required the import of
`plumpy` which has a non-negligible import time. This is a breaking
change but it is not expected to be used outside of `aiida-core`.
  • Loading branch information
danielhollas committed Sep 5, 2023
1 parent ae637d8 commit 8e6e08d
Show file tree
Hide file tree
Showing 17 changed files with 47 additions and 70 deletions.
37 changes: 0 additions & 37 deletions .github/system_tests/test_verdi_load_time.sh

This file was deleted.

1 change: 1 addition & 0 deletions .github/workflows/ci-code.yml
Original file line number Diff line number Diff line change
Expand Up @@ -148,4 +148,5 @@ jobs:
- name: Run verdi
run: |
verdi devel check-load-time
verdi devel check-undesired-imports
.github/workflows/verdi.sh
6 changes: 1 addition & 5 deletions .github/workflows/verdi.sh
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ VERDI=`which verdi`
# tends to go towards ~0.8 seconds. Since these timings are obviously machine and environment dependent, typically these
# types of tests are fragile. But with a load limit of more than twice the ideal loading time, if exceeded, should give
# a reasonably sure indication that the loading of `verdi` is unacceptably slowed down.
LOAD_LIMIT=0.5
LOAD_LIMIT=0.4
MAX_NUMBER_ATTEMPTS=5

iteration=0
Expand All @@ -35,10 +35,6 @@ while true; do

done

$VERDI devel check-load-time
$VERDI devel check-undesired-imports


# Test that we can also run the CLI via `python -m aiida`,
# that it returns a 0 exit code, and contains the expected stdout.
echo "Invoking verdi via `python -m aiida`"
Expand Down
5 changes: 4 additions & 1 deletion aiida/cmdline/commands/cmd_devel.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,10 @@ def devel_check_undesired_imports():
"""
loaded_modules = 0

for modulename in ['seekpath', 'CifFile', 'ase', 'pymatgen', 'spglib', 'pymysql']:
for modulename in [
'asyncio', 'requests', 'plumpy', 'disk_objectstore', 'paramiko', 'seekpath', 'CifFile', 'ase', 'pymatgen',
'spglib', 'pymysql'
]:
if modulename in sys.modules:
echo.echo_warning(f'Detected loaded module "{modulename}"')
loaded_modules += 1
Expand Down
5 changes: 3 additions & 2 deletions aiida/cmdline/commands/cmd_rabbitmq.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
import typing as t

import click
import requests
import tabulate
import wrapt
import yaml
Expand All @@ -27,6 +26,7 @@

if t.TYPE_CHECKING:
import kiwipy.rmq
import requests

from aiida.manage.configuration.profile import Profile

Expand Down Expand Up @@ -91,12 +91,13 @@ def cmd_tasks():
)


def echo_response(response: requests.Response, exit_on_error: bool = True) -> None:
def echo_response(response: 'requests.Response', exit_on_error: bool = True) -> None:
"""Echo the response of a request.
:param response: The response to the request.
:param exit_on_error: Boolean, if ``True``, call ``sys.exit`` with the status code of the response.
"""
import requests
try:
response.raise_for_status()
except requests.HTTPError:
Expand Down
1 change: 0 additions & 1 deletion aiida/manage/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,6 @@
'Option',
'Postgres',
'PostgresConnectionMode',
'ProcessLauncher',
'Profile',
'RabbitmqManagementClient',
'check_and_migrate_config',
Expand Down
3 changes: 1 addition & 2 deletions aiida/manage/configuration/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,6 @@
import tempfile
from typing import Any, Dict, Optional, Sequence, Tuple

import jsonschema

from aiida.common.exceptions import ConfigurationError

from . import schema as schema_module
Expand Down Expand Up @@ -126,6 +124,7 @@ def _backup(cls, filepath):
@staticmethod
def validate(config: dict, filepath: Optional[str] = None):
"""Validate a configuration dictionary."""
import jsonschema
try:
jsonschema.validate(instance=config, schema=config_schema())
except jsonschema.ValidationError as error:
Expand Down
4 changes: 2 additions & 2 deletions aiida/manage/configuration/options.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,6 @@
"""Definition of known configuration options and methods to parse and get option values."""
from typing import Any, Dict, List, Tuple

import jsonschema

from aiida.common.exceptions import ConfigurationError

__all__ = ('get_option', 'get_option_names', 'parse_option', 'Option')
Expand Down Expand Up @@ -64,6 +62,8 @@ def validate(self, value: Any, cast: bool = True) -> Any:
"""
# pylint: disable=too-many-branches
import jsonschema

from aiida.manage.caching import _validate_identifier_pattern

from .config import ConfigValidationError
Expand Down
1 change: 0 additions & 1 deletion aiida/manage/external/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@
'ManagementApiConnectionError',
'Postgres',
'PostgresConnectionMode',
'ProcessLauncher',
'RabbitmqManagementClient',
'get_launch_queue_name',
'get_message_exchange_name',
Expand Down
2 changes: 0 additions & 2 deletions aiida/manage/external/rmq/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,13 +16,11 @@

from .client import *
from .defaults import *
from .launcher import *
from .utils import *

__all__ = (
'BROKER_DEFAULTS',
'ManagementApiConnectionError',
'ProcessLauncher',
'RabbitmqManagementClient',
'get_launch_queue_name',
'get_message_exchange_name',
Expand Down
9 changes: 6 additions & 3 deletions aiida/manage/external/rmq/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,11 @@
import typing as t
from urllib.parse import quote

import requests

from aiida.common.exceptions import AiidaException

if t.TYPE_CHECKING:
import requests

__all__ = ('RabbitmqManagementClient', 'ManagementApiConnectionError')


Expand All @@ -31,6 +32,7 @@ def __init__(self, username: str, password: str, hostname: str, virtual_host: st
:param hostname: The hostname of the RabbitMQ server.
:param virtual_host: The virtual host.
"""
import requests
self._username = username
self._password = password
self._hostname = hostname
Expand Down Expand Up @@ -58,7 +60,7 @@ def request(
url_params: dict[str, str] | None = None,
method: str = 'GET',
params: dict[str, t.Any] | None = None,
) -> requests.Response:
) -> 'requests.Response':
"""Make a request.
:param url: The resource path with placeholders, e.g., ``queues/{virtual_host}/{queue}``.
Expand All @@ -69,6 +71,7 @@ def request(
:returns: The response of the request.
:raises `ManagementApiConnectionError`: If connection to the API cannot be made.
"""
import requests
url = self.format_url(url, url_params)
try:
return requests.request(method, url, auth=self._authentication, params=params or {}, timeout=5)
Expand Down
2 changes: 0 additions & 2 deletions aiida/manage/external/rmq/launcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,6 @@

LOGGER = logging.getLogger(__name__)

__all__ = ('ProcessLauncher',)


class ProcessLauncher(plumpy.ProcessLauncher):
"""A sub class of :class:`plumpy.ProcessLauncher` to launch a ``Process``.
Expand Down
9 changes: 5 additions & 4 deletions aiida/manage/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,11 +9,12 @@
###########################################################################
# pylint: disable=cyclic-import
"""AiiDA manager for global settings"""
import asyncio
import functools
from typing import TYPE_CHECKING, Any, Optional, Union

if TYPE_CHECKING:
import asyncio

from kiwipy.rmq import RmqThreadCommunicator
from plumpy.process_comms import RemoteProcessThreadController

Expand Down Expand Up @@ -416,7 +417,7 @@ def create_runner(self, with_persistence: bool = True, **kwargs: Any) -> 'Runner

return runners.Runner(**settings)

def create_daemon_runner(self, loop: Optional[asyncio.AbstractEventLoop] = None) -> 'Runner':
def create_daemon_runner(self, loop: Optional['asyncio.AbstractEventLoop'] = None) -> 'Runner':
"""Create and return a new daemon runner.
This is used by workers when the daemon is running and in testing.
Expand All @@ -429,13 +430,13 @@ def create_daemon_runner(self, loop: Optional[asyncio.AbstractEventLoop] = None)
from plumpy.persistence import LoadSaveContext

from aiida.engine import persistence
from aiida.manage.external import rmq
from aiida.manage.external.rmq.launcher import ProcessLauncher

runner = self.create_runner(rmq_submit=True, loop=loop)
runner_loop = runner.loop

# Listen for incoming launch requests
task_receiver = rmq.ProcessLauncher(
task_receiver = ProcessLauncher(
loop=runner_loop,
persister=self.get_persister(),
load_context=LoadSaveContext(runner=runner),
Expand Down
9 changes: 6 additions & 3 deletions aiida/repository/backend/disk_object_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,13 +4,14 @@
import shutil
import typing as t

from disk_objectstore import Container

from aiida.common.lang import type_check
from aiida.storage.log import STORAGE_LOGGER

from .abstract import AbstractRepositoryBackend

if t.TYPE_CHECKING:
from disk_objectstore import Container # pylint: disable=unused-import

__all__ = ('DiskObjectStoreRepositoryBackend',)

BYTES_TO_MB = 1 / 1024**2
Expand All @@ -30,7 +31,9 @@ class DiskObjectStoreRepositoryBackend(AbstractRepositoryBackend):
"""

def __init__(self, container: Container):
def __init__(self, container: 'Container'):
if not t.TYPE_CHECKING:
from disk_objectstore import Container # pylint: disable=redefined-outer-name
type_check(container, Container)
self._container = container

Expand Down
6 changes: 5 additions & 1 deletion aiida/storage/psql_dos/backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
import pathlib
from typing import TYPE_CHECKING, Iterator, List, Optional, Sequence, Set, Union

from disk_objectstore import Container
from sqlalchemy.orm import Session, scoped_session, sessionmaker

from aiida.common.exceptions import ClosedStorage, ConfigurationError, IntegrityError
Expand All @@ -29,6 +28,8 @@
from .orm import authinfos, comments, computers, convert, groups, logs, nodes, querybuilder, users

if TYPE_CHECKING:
from disk_objectstore import Container

from aiida.repository.backend import DiskObjectStoreRepositoryBackend

__all__ = ('PsqlDosBackend',)
Expand Down Expand Up @@ -195,7 +196,10 @@ def _clear(self) -> None:
)

def get_repository(self) -> 'DiskObjectStoreRepositoryBackend':
from disk_objectstore import Container

from aiida.repository.backend import DiskObjectStoreRepositoryBackend

container = Container(str(get_filepath_container(self.profile)))
return DiskObjectStoreRepositoryBackend(container=container)

Expand Down
11 changes: 8 additions & 3 deletions aiida/storage/psql_dos/migrator.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,14 +19,13 @@

import contextlib
import pathlib
from typing import Any, Dict, Iterator, Optional
from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional

from alembic.command import downgrade, upgrade
from alembic.config import Config
from alembic.runtime.environment import EnvironmentContext
from alembic.runtime.migration import MigrationContext, MigrationInfo
from alembic.script import ScriptDirectory
from disk_objectstore import Container
from sqlalchemy import MetaData, String, column, desc, insert, inspect, select, table
from sqlalchemy.exc import OperationalError, ProgrammingError
from sqlalchemy.ext.automap import automap_base
Expand All @@ -38,6 +37,9 @@
from aiida.storage.psql_dos.models.settings import DbSetting
from aiida.storage.psql_dos.utils import create_sqlalchemy_engine

if TYPE_CHECKING:
from disk_objectstore import Container

TEMPLATE_LEGACY_DJANGO_SCHEMA = """
Database schema is using the legacy Django schema.
To migrate the database schema version to the current one, run the following command:
Expand Down Expand Up @@ -175,12 +177,15 @@ def validate_storage(self) -> None:
f'but the disk-objectstore\'s is {repository_uuid}.'
)

def get_container(self) -> Container:
def get_container(self) -> 'Container':
"""Return the disk-object store container.
:returns: The disk-object store container configured for the repository path of the current profile.
"""
from disk_objectstore import Container

from .backend import get_filepath_container

return Container(get_filepath_container(self.profile))

def get_repository_uuid(self) -> str:
Expand Down
6 changes: 5 additions & 1 deletion aiida/transports/plugins/ssh.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
from stat import S_ISDIR, S_ISREG

import click
import paramiko

from aiida.cmdline.params import options
from aiida.cmdline.params.types.path import AbsolutePathOrEmptyParamType
Expand All @@ -36,6 +35,8 @@ def parse_sshconfig(computername):
:param computername: the computer name for which we want the configuration.
"""
import paramiko

config = paramiko.SSHConfig()
try:
with open(os.path.expanduser('~/.ssh/config'), encoding='utf8') as fhandle:
Expand Down Expand Up @@ -397,6 +398,8 @@ def __init__(self, *args, **kwargs):
function (as port, username, password, ...); taken from the
accepted paramiko.SSHClient.connect() params.
"""
import paramiko

super().__init__(*args, **kwargs)

self._sftp = None
Expand Down Expand Up @@ -440,6 +443,7 @@ def open(self): # pylint: disable=too-many-branches,too-many-statements
:raise aiida.common.InvalidOperation: if the channel is already open
"""
import paramiko
from paramiko.ssh_exception import SSHException

from aiida.common.exceptions import InvalidOperation
Expand Down

0 comments on commit 8e6e08d

Please sign in to comment.