gigl.utils.share_memory

Attributes

logger

Functions

allocate_disk_backed(shape, dtype)	An EMPTY writable tensor whose bytes live in a file, or None to fall back to RAM.
allocate_preshared(shape, dtype[, random_access])	Allocate a large tensor in its FINAL home, so share_memory_() cannot duplicate it.
has_live_mapping(path)	Whether THIS process still maps path, read from /proc/self/maps.
is_disk_backed(tensor)	Whether tensor (or a view of one) is an mmap over a real file rather than RAM.
is_tensor_spilling_enabled()	Whether GIGL_TENSOR_SPILL_DIR is set, i.e. large tensors should go to disk.
load_spilled_tensor(path, dtype, shape)	Re-map a spilled tensor in the current process, without copying its bytes anywhere.
prepare_spill_dir()	Clear leftover spill files, ONCE per run, before anything in the run can spill.
release_page_cache(tensor)	Write a disk-backed tensor's dirty pages out and ask the kernel to drop them from the cache.
release_page_cache_by_path(path)	Request eviction of a spill file's page cache when nothing maps it any more.
share_memory(entity)	Based on GraphLearn-for-PyTorch's share_memory implementation, with additional support for handling empty tensors with share_memory.
share_memory_for_ipc(entity)	Prepare a mapping of tensors to cross a process boundary, spilling instead of copying.
spill_tensor_to_disk(tensor)	Write tensor to disk and return a tensor mapped over the file, or None to keep it.

Module Contents

gigl.utils.share_memory.allocate_disk_backed(shape, dtype)

An EMPTY writable tensor whose bytes live in a file, or None to fall back to RAM.

The counterpart to spill_tensor_to_disk() for a buffer that does not exist yet. Spilling fills RAM and then copies it out, so it needs the full size in anonymous memory at least once; allocating the destination as the file from the start means those bytes never occupy anonymous memory at all.

Blocks are RESERVED up front with posix_fallocate rather than left sparse, and reservation is mandatory. from_file gives the file its full apparent size without allocating it, so a filesystem that fills later fails at the moment a page is written – and a write to a mapping that cannot be backed raises SIGBUS, which is a signal, not an exception: it bypasses every try here and kills the process with no traceback. Reserving turns that into an OSError at allocation time, where it can be reported and fallen back on; a filesystem that cannot reserve at all gets None rather than an unreserved mapping.

Sequential writes through the mapping cost ~3.1x their in-memory equivalent (measured on a wide fp32 matrix written row-by-row). What that buys: the destination’s bytes stop being unreclaimable anonymous memory – still charged to the cgroup while resident, but reclaimable page cache the kernel can evict under pressure instead of OOM-killing.

Returns None whenever a file-backed buffer is not available or not worth it – spilling disabled, below the size threshold, quantized dtype, no room to reserve, reservation unsupported by the filesystem, IO failure – so callers can simply fall back to torch.empty.

Parameters:

• shape (tuple[int, Ellipsis])

• dtype (torch.dtype)

Return type:

Optional[torch.Tensor]

gigl.utils.share_memory.allocate_preshared(shape, dtype, random_access=False)

Allocate a large tensor in its FINAL home, so share_memory_() cannot duplicate it.

GLT’s Graph.__init__ shares the topology unconditionally, and for an anonymous tensor that copies every byte into /dev/shm – the tensor exists twice during the copy, fatal for a tens-of-GiB CSR near the memory limit. Both homes here are immune to that copy: a spill file (consumers that check is_disk_backed() leave it alone), or POSIX shared memory allocated directly (share_memory_() then finds it already shared and does nothing).

Disk-first for a destination written in order. A scatter destination prefers memory: an 8-byte write to an uncached file page is a 4 KiB read-modify-write, repeated over the same pages, so a file-backed scatter destination behaves like a hung run. With random_access=True a file is only the checked, loudly-logged fallback for when memory will not hold the tensor.

Below the spill threshold returns a plain tensor. Uninitialised in all cases, like torch.empty.

Parameters:

• shape (tuple[int, Ellipsis]) – Shape of the tensor.

• dtype (torch.dtype) – Dtype of the tensor.

• random_access (bool) – Set when the caller will write to scattered offsets rather than stream through in order; memory is then preferred and disk the checked fallback.

Return type:

torch.Tensor

gigl.utils.share_memory.has_live_mapping(path)

Whether THIS process still maps path, read from /proc/self/maps.

The kernel’s own list of this process’s mappings, so it sees every mapping over the file no matter who created it. False when the list cannot be read (non-Linux), matching the rest of this module: spilling is effectively Linux-only and refuses gracefully elsewhere.

Parameters:

path (str)

Return type:

bool

gigl.utils.share_memory.is_disk_backed(tensor)

Whether tensor (or a view of one) is an mmap over a real file rather than RAM.

Read off the tensor itself: a storage created by from_file(shared=True) records the path in untyped_storage().filename, and a view shares its base’s storage, so the label travels with it. filename is None for everything else – including tmpfs-shared storages (share_memory_(), _new_shared), which are RAM and must not be treated as spilled.

Parameters:

tensor (torch.Tensor)

Return type:

bool

gigl.utils.share_memory.is_tensor_spilling_enabled()

Whether GIGL_TENSOR_SPILL_DIR is set, i.e. large tensors should go to disk.

Return type:

bool

gigl.utils.share_memory.load_spilled_tensor(path, dtype, shape)

Re-map a spilled tensor in the current process, without copying its bytes anywhere.

The counterpart to spill_tensor_to_disk() for a path that arrived through some channel other than pickling a tensor (pickling already re-maps by itself).

Raises:

• ValueError – If the file is the wrong size for shape – which would otherwise be read as silently wrong data.

• OSError – If the file cannot be mapped writable.

Parameters:

• path (str)

• dtype (torch.dtype)

• shape (tuple[int, Ellipsis])

Return type:

torch.Tensor

gigl.utils.share_memory.prepare_spill_dir()

Clear leftover spill files, ONCE per run, before anything in the run can spill.

Cleanup is done at the START of a run, not with atexit. The tensors are spilled inside the mp.spawn dataset-building child (dataset_factory._build_dataset_process), and that child exits before the trainer uses the dataset – so unlinking on its exit would delete files the trainer and its sampling workers still have mapped by path. Removing stale files up front bounds disk use without that hazard.

Exactly one process may do this, and it must be a process that starts before any spilling one – a run has several spilling children in sequence, and a later one cannot tell a sibling’s live files from a stale run’s (the first spiller does clean up itself when no ancestor did, which covers standalone use and tests). Per-process age-based cleanup is NOT a substitute and is actively wrong: with sequential loading, the edge child spills and exits, then the node child starts, sees the edge files as older than itself, and deletes files the parent has not mapped yet. The marker below is an environment variable precisely because spawn children inherit the environment, so a child can tell that its parent already did this.

Idempotent, and safe to call when spilling is disabled.

Return type:

None

gigl.utils.share_memory.release_page_cache(tensor)

Write a disk-backed tensor’s dirty pages out and ask the kernel to drop them from the cache.

A spilled tensor’s pages are reclaimable but still charged to cgroup v2’s memory.current while resident; dropping them releases the charge. The mapping stays valid and refaults on next access, so this is worth it only for data written now and not read until much later. Operates on the tensor’s whole backing file, so views over the same file lose their cached pages too.

Three steps, in this order, and every one is load-bearing:

1. fsync – writes back the pages dirtied through the mapping; dirty pages cannot be dropped at all.

2. MADV_DONTNEED on the mapping – removes this process’s page-table entries. Without it step 3 silently skips every still-mapped page and reports success having freed nothing.

3. POSIX_FADV_DONTNEED – discards the now-unmapped clean page cache.

Returns True when every step COMPLETED – FADV_DONTNEED is advisory, so completion is not proof of eviction; measure residency with mincore to know (the tests do). False means the sequence was incomplete, though a failure after step 2 has already dropped this process’s page-table entries.

Parameters:

tensor (torch.Tensor)

Return type:

bool

gigl.utils.share_memory.release_page_cache_by_path(path)

Request eviction of a spill file’s page cache when nothing maps it any more.

The counterpart to release_page_cache() for a tensor already dropped: del unmaps it but leaves the file’s pages charged to the cgroup until the kernel reclaims them.

Refuses when a mapping is still live: FADV_DONTNEED silently skips pages present in any page table, so it would return success having freed nothing. Use release_page_cache() on the tensor in that case, which unmaps first.

Returns True if the eviction was REQUESTED (FADV_DONTNEED is advisory and reports no count); False if a mapping is still live, the file could not be opened, or the platform lacks posix_fadvise.

Parameters:

path (str)

Return type:

bool

gigl.utils.share_memory.share_memory(entity)

Based on GraphLearn-for-PyTorch’s share_memory implementation, with additional support for handling empty tensors with share_memory.

alibaba/graphlearn-for-pytorch

Calling share_memory_() on an empty tensor may cause processes to hang, although the root cause of this is currently unknown. As a result, we opt to not move empty tensors to shared memory if they are provided.

When calling share_memory on a RangePartitionBook, we don’t need to move the partition bounds to shared memory, since GLT doesn’t natively provide a ForkingPickler registration method for the RangePartitionBook, and the cost of not moving this to shared memory is minimal, since the size of this tensor is very small, being equal in length to the number of machines.

This function never spills to disk, and a tensor that is ALREADY disk-backed is left alone: share_memory_() on an mmap-backed tensor copies every byte into /dev/shm, quietly undoing a spill made elsewhere. Spilling belongs to share_memory_for_ipc() and spill_tensor_to_disk().

Parameters:

entity (Optional[Union[torch.Tensor, PartitionBook, dict[_KeyType, torch.Tensor], dict[_KeyType, PartitionBook]]]) – Homogeneous or heterogeneous entity of tensors which is being moved to shared memory

Return type:

None

gigl.utils.share_memory.share_memory_for_ipc(entity)

Prepare a mapping of tensors to cross a process boundary, spilling instead of copying.

Same intent as share_memory(), but for values that will be pickled to another process: each value large enough is spilled first, and a spilled tensor pickles as its file path (see _reduce_spill_aware()), so its bytes never transit /dev/shm. Tensors that are not spilled (spilling disabled, below the size threshold) go to POSIX shared memory exactly as before, which pickles by handle already.

Returns a NEW dict; the input is left alone, since the caller usually still holds the tensors and dropping them is its decision.

Parameters:

entity (dict[_KeyType, torch.Tensor])

Return type:

dict[_KeyType, torch.Tensor]

gigl.utils.share_memory.spill_tensor_to_disk(tensor)

Write tensor to disk and return a tensor mapped over the file, or None to keep it.

None covers every reason not to spill – spilling disabled, tensor below the threshold, IO failure – so callers can treat it as “keep what you had”. The returned tensor’s storage knows its file (is_disk_backed()), so share_memory() leaves it alone and pickling ships the path rather than the bytes.

Exposed for callers that hold a tensor directly rather than inside a Mapping, which share_memory() cannot substitute into. The partitioned node feature matrix is the case that matters: it is created after loading, so the spill in load_torch_tensors never sees it, and it is resident during the graph build.

Parameters:

tensor (torch.Tensor)

Return type:

Optional[torch.Tensor]

gigl.utils.share_memory.logger