Manual memory pinning

packages/common/pyproject.toml

@@ -35,7 +35,10 @@ missing-attribute = false
 no-matching-overload = false
 bad-context-manager = false
 
-
+[tool.pyrefly.ignores]
+import-error = [
+    "torch",
+]
 
 
 # The linting configuration

packages/common/src/weathergen/common/io.py

@@ -16,6 +16,7 @@
 
 import dask.array as da
 import numpy as np
+import torch
 import xarray as xr
 import zarr
 from numpy import datetime64
@@ -109,6 +110,16 @@ def is_empty(self):
         """
         return len(self.data) == 0
 
+    def pin_memory(self):
+        """Pin all tensors in IOReaderData"""
+        if hasattr(self, "coords") and isinstance(self.coords, torch.Tensor):
+            self.coords = self.coords.pin_memory()
+        if hasattr(self, "data") and isinstance(self.data, torch.Tensor):
+            self.data = self.data.pin_memory()
+        if hasattr(self, "geoinfos") and isinstance(self.geoinfos, torch.Tensor):
+            self.geoinfos = self.geoinfos.pin_memory()
+        return self
+
     @classmethod
     def create(cls, other: typing.Any) -> "IOReaderData":
         """

packages/evaluate/src/weathergen/evaluate/export/export_inference.py

@@ -155,7 +155,7 @@ def parse_args(args: list) -> argparse.Namespace:
         type=str,
         help="Grid type to include in the output filename (i.e. 'O96/N320')",
         required=False,
-        default="O96", 
+        default="O96",
         dest="quaver_template_grid_type",
     )
 

src/weathergen/datasets/batch.py

@@ -37,6 +37,25 @@ class Sample:
     # keys: stream_name, values: StreamData
     streams_data: dict[str, StreamData | None]
 
+    def pin_memory(self):
+        """Pin all tensors in this Sample to CPU pinned memory"""
+
+        # Pin StreamData objects in streams_data dict
+        if hasattr(self, "streams_data") and isinstance(self.streams_data, dict):
+            for _stream_name, stream_data in self.streams_data.items():
+                if stream_data is not None and hasattr(stream_data, "pin_memory"):
+                    stream_data.pin_memory()
+
+        # Pin tensors in meta_info
+        if hasattr(self, "meta_info") and isinstance(self.meta_info, dict):
+            for _key, meta_data in self.meta_info.items():
+                if isinstance(meta_data, SampleMetaData):
+                    # Pin mask tensor
+                    if meta_data.mask is not None and isinstance(meta_data.mask, torch.Tensor):
+                        meta_data.mask = meta_data.mask.pin_memory()
+
+        return self
+
     def __init__(self, streams: dict) -> None:
         # TODO: can we pass this right away?
         self.meta_info = {}
@@ -124,6 +143,23 @@ def __init__(self, streams, num_source_samples: int, num_target_samples: int) ->
         self.source2target_matching_idxs = np.full(num_source_samples, -1, dtype=np.int32)
         self.target2source_matching_idxs = [[] for _ in range(num_target_samples)]
 
+    def pin_memory(self):
+        """Pin all tensors in this batch to CPU pinned memory"""
+
+        # Pin all source samples
+        for sample in self.source_samples:
+            sample.pin_memory()
+
+        # Pin all target samples
+        for sample in self.target_samples:
+            sample.pin_memory()
+
+        # Pin source_tokens_lens
+        if isinstance(self.source_tokens_lens, torch.Tensor):
+            self.source_tokens_lens = self.source_tokens_lens.pin_memory()
+
+        return self
+
     def to_device(self, device):  # -> ModelBatch
         for sample in self.source_samples:
             sample.to_device(device)

src/weathergen/datasets/multi_stream_data_sampler.py

@@ -762,6 +762,9 @@ def __iter__(self) -> ModelBatch:
 
                 batch = self._get_batch(idx, forecast_dt)
 
+                # pin memory for faster CPU-GPU transfer
+                batch = batch.pin_memory()
+
                 # skip completely empty batch item or when all targets are empty -> no grad
                 if not batch.is_empty():
                     break

src/weathergen/datasets/stream_data.py

@@ -75,6 +75,63 @@ def __init__(self, idx: int, input_steps: int, forecast_steps: int, healpix_cell
         self.source_idxs_embed = [torch.tensor([]) for _ in range(self.input_steps)]
         self.source_idxs_embed_pe = [torch.tensor([]) for _ in range(self.input_steps)]
 
+    def pin_memory(self):
+        """Pin all tensors in this StreamData object to CPU pinned memory"""
+
+        # Pin target tensors
+        self.target_coords = [
+            t.pin_memory() if isinstance(t, torch.Tensor) and t.numel() > 0 else t
+            for t in self.target_coords
+        ]
+        self.target_coords_lens = [
+            t.pin_memory() if isinstance(t, torch.Tensor) and t.numel() > 0 else t
+            for t in self.target_coords_lens
+        ]
+        self.target_tokens = [
+            t.pin_memory() if isinstance(t, torch.Tensor) and t.numel() > 0 else t
+            for t in self.target_tokens
+        ]
+        self.target_tokens_lens = [
+            t.pin_memory() if isinstance(t, torch.Tensor) and t.numel() > 0 else t
+            for t in self.target_tokens_lens
+        ]
+        self.idxs_inv = [
+            t.pin_memory() if isinstance(t, torch.Tensor) and t.numel() > 0 else t
+            for t in self.idxs_inv
+        ]
+
+        # Pin target_coords_raw (list of tensors)
+        self.target_coords_raw = [
+            t.pin_memory() if isinstance(t, torch.Tensor) and t.numel() > 0 else t
+            for t in self.target_coords_raw
+        ]
+
+        # Pin source tensors
+        self.source_tokens_cells = [
+            s.pin_memory() if s is not None and isinstance(s, torch.Tensor) else s
+            for s in self.source_tokens_cells
+        ]
+        self.source_tokens_lens = [
+            s.pin_memory() if isinstance(s, torch.Tensor) and s.numel() > 0 else s
+            for s in self.source_tokens_lens
+        ]
+        self.source_idxs_embed = [
+            s.pin_memory() if isinstance(s, torch.Tensor) and s.numel() > 0 else s
+            for s in self.source_idxs_embed
+        ]
+        self.source_idxs_embed_pe = [
+            s.pin_memory() if isinstance(s, torch.Tensor) and s.numel() > 0 else s
+            for s in self.source_idxs_embed_pe
+        ]
+
+        # Pin source_raw (list of IOReaderData objects)
+        if hasattr(self, "source_raw"):
+            for raw_data in self.source_raw:
+                if raw_data is not None and hasattr(raw_data, "pin_memory"):
+                    raw_data.pin_memory()
+
+        return self
+
     def to_device(self, device: str) -> None:
         """
         Move data to GPU