Changes as of comments from 6 Mar 2024, from PR 14

ndsl/__init__.py

@@ -1,9 +1,7 @@
-from .buffer import Buffer
-from .comm.boundary import Boundary, SimpleBoundary
 from .comm.communicator import CubedSphereCommunicator, TileCommunicator
-from .comm.local_comm import AsyncResult, ConcurrencyError, LocalComm
+from .comm.local_comm import LocalComm
 from .comm.mpi import MPIComm
-from .comm.null_comm import NullAsyncResult, NullComm
+from .comm.null_comm import NullComm
 from .comm.partitioner import CubedSpherePartitioner, TilePartitioner
 from .constants import ConstantVersions
 from .dsl.caches.codepath import FV3CodePath
@@ -16,46 +14,24 @@
     StorageReport,
 )
 from .dsl.dace.wrapped_halo_exchange import WrappedHaloUpdater
-from .dsl.stencil import (
-    CompareToNumpyStencil,
-    FrozenStencil,
-    GridIndexing,
-    StencilFactory,
-    TimingCollector,
-)
+from .dsl.stencil import FrozenStencil, GridIndexing, StencilFactory, TimingCollector
 from .dsl.stencil_config import CompilationConfig, RunMode, StencilConfig
 from .exceptions import OutOfBoundsError
-from .halo.data_transformer import (
-    HaloDataTransformer,
-    HaloDataTransformerCPU,
-    HaloDataTransformerGPU,
-    HaloExchangeSpec,
-)
+from .halo.data_transformer import HaloExchangeSpec
 from .halo.updater import HaloUpdater, HaloUpdateRequest, VectorInterfaceHaloUpdater
-from .initialization.allocator import QuantityFactory, StorageNumpy
+from .initialization.allocator import QuantityFactory
 from .initialization.sizer import GridSizer, SubtileGridSizer
 from .logging import ndsl_log
 from .monitor.netcdf_monitor import NetCDFMonitor
 from .monitor.protocol import Protocol
 from .monitor.zarr_monitor import ZarrMonitor
 from .namelist import Namelist
-from .optional_imports import RaiseWhenAccessed
-from .performance.collector import (
-    AbstractPerformanceCollector,
-    NullPerformanceCollector,
-    PerformanceCollector,
-)
+from .performance.collector import NullPerformanceCollector, PerformanceCollector
 from .performance.config import PerformanceConfig
 from .performance.profiler import NullProfiler, Profiler
 from .performance.report import Experiment, Report, TimeReport
 from .performance.timer import NullTimer, Timer
-from .quantity import (
-    BoundaryArrayView,
-    BoundedArrayView,
-    Quantity,
-    QuantityHaloSpec,
-    QuantityMetadata,
-)
+from .quantity import Quantity
 from .testing.dummy_comm import DummyComm
 from .types import Allocator, AsyncRequest, NumpyModule
 from .units import UnitsError

ndsl/checkpointer/base.py

@@ -0,0 +1,7 @@
+import abc
+
+
+class Checkpointer(abc.ABC):
+    @abc.abstractmethod
+    def __call__(self, savepoint_name, **kwargs):
+        ...

ndsl/checkpointer/null.py

@@ -1,4 +1,4 @@
-from ndsl.typing import Checkpointer
+from ndsl.checkpointer.base import Checkpointer
 
 
 class NullCheckpointer(Checkpointer):

ndsl/checkpointer/snapshots.py

@@ -2,9 +2,9 @@
 
 import numpy as np
 
+from ndsl.checkpointer.base import Checkpointer
 from ndsl.optional_imports import cupy as cp
 from ndsl.optional_imports import xarray as xr
-from ndsl.typing import Checkpointer
 
 
 def make_dims(savepoint_dim, label, data_list):

ndsl/checkpointer/thresholds.py

@@ -5,8 +5,8 @@
 
 import numpy as np
 
+from ndsl.checkpointer.base import Checkpointer
 from ndsl.quantity import Quantity
-from ndsl.typing import Checkpointer
 
 
 try:

ndsl/checkpointer/validation.py

@@ -5,14 +5,14 @@
 
 import numpy as np
 
+from ndsl.checkpointer.base import Checkpointer
 from ndsl.checkpointer.thresholds import (
     ArrayLike,
     SavepointName,
     SavepointThresholds,
     cast_to_ndarray,
 )
 from ndsl.optional_imports import xarray as xr
-from ndsl.typing import Checkpointer
 
 
 def _clip_pace_array_to_target(