Skip to content

Commit

Permalink
PERF: delay rarely used imports in yt.utilities
Browse files Browse the repository at this point in the history
  • Loading branch information
neutrinoceros committed Jun 27, 2023
1 parent 5d070c3 commit 2c52296
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 4 deletions.
14 changes: 11 additions & 3 deletions yt/utilities/parallel_tools/parallel_analysis_interface.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,7 @@
import numpy as np
from more_itertools import always_iterable

import yt.utilities.logger
from yt.config import ytcfg
from yt.data_objects.image_array import ImageArray
from yt.funcs import is_sequence
from yt.units.unit_registry import UnitRegistry # type: ignore
from yt.units.yt_array import YTArray
Expand Down Expand Up @@ -63,7 +61,7 @@ def write_to_file(exc_type, exc, tb):
def default_mpi_excepthook(exception_type, exception_value, tb):
traceback.print_tb(tb)
mylog.error("%s: %s", exception_type.__name__, exception_value)
comm = yt.communication_system.communicators[-1]
comm = communication_system.communicators[-1]
if comm.size > 1:
mylog.error("Error occurred on rank %d.", comm.rank)
MPI.COMM_WORLD.Abort(1)
Expand All @@ -90,6 +88,8 @@ def enable_parallelism(suppress_logging: bool = False, communicator=None) -> boo
parallel_capable: bool
True if the call was successful. False otherwise.
"""
import yt.utilities.logger

global parallel_capable, MPI
try:
from mpi4py import MPI as _MPI
Expand Down Expand Up @@ -828,6 +828,8 @@ def par_combine_object(self, data, op, datatype=None):

@parallel_passthrough
def mpi_bcast(self, data, root=0):
from yt.data_objects.image_array import ImageArray

# The second check below makes sure that we know how to communicate
# this type of array. Otherwise, we'll pickle it.
if isinstance(data, np.ndarray) and get_mpi_type(data.dtype) is not None:
Expand Down Expand Up @@ -1076,6 +1078,8 @@ def merge_quadtree_buffers(self, qt, merge_style):
return qt

def send_array(self, arr, dest, tag=0):
from yt.data_objects.image_array import ImageArray

if not isinstance(arr, np.ndarray):
self.comm.send((None, None), dest=dest, tag=tag)
self.comm.send(arr, dest=dest, tag=tag)
Expand All @@ -1095,6 +1099,8 @@ def send_array(self, arr, dest, tag=0):
del tmp

def recv_array(self, source, tag=0):
from yt.data_objects.image_array import ImageArray

metadata = self.comm.recv(source=source, tag=tag)
dt, ne = metadata[:2]
if ne is None and dt is None:
Expand All @@ -1111,6 +1117,8 @@ def recv_array(self, source, tag=0):
return arr

def alltoallv_array(self, send, total_size, offsets, sizes):
from yt.data_objects.image_array import ImageArray

if len(send.shape) > 1:
recv = []
for i in range(send.shape[0]):
Expand Down
5 changes: 4 additions & 1 deletion yt/utilities/parameter_file_storage.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
import csv
import os.path
from itertools import islice

Expand Down Expand Up @@ -170,6 +169,8 @@ def get_recent(self, n=10):

@parallel_simple_proxy
def _write_out(self):
import csv

if self._read_only:
return
fn = self._get_db_name()
Expand All @@ -187,6 +188,8 @@ def _write_out(self):
@parallel_simple_proxy
def read_db(self):
"""This will read the storage device from disk."""
import csv

f = open(self._get_db_name())
vals = csv.DictReader(f, _field_names)
db = {}
Expand Down

0 comments on commit 2c52296

Please sign in to comment.