From 4c5871d9b9c7cf33dbbf31a0f9d82b5175309dff Mon Sep 17 00:00:00 2001 From: afonso pinto Date: Sun, 10 May 2026 22:09:48 +0100 Subject: [PATCH] feat: Improve GPU memory usage indication --- src/ui/layer_bar.css | 16 +++- src/ui/layer_bar.ts | 64 ++++++++++++++- src/ui/layer_bar_gpu_memory_pressure.spec.ts | 67 +++++++++++++++ src/ui/layer_bar_gpu_memory_pressure.ts | 86 ++++++++++++++++++++ 4 files changed, 231 insertions(+), 2 deletions(-) create mode 100644 src/ui/layer_bar_gpu_memory_pressure.spec.ts create mode 100644 src/ui/layer_bar_gpu_memory_pressure.ts diff --git a/src/ui/layer_bar.css b/src/ui/layer_bar.css index d9c7877f06..aacff3dc86 100644 --- a/src/ui/layer_bar.css +++ b/src/ui/layer_bar.css @@ -112,10 +112,24 @@ .neuroglancer-layer-item-prefetch-progress { position: absolute; left: 0px; - height: 2px; + height: 4px; background-color: #666; } +.neuroglancer-layer-item[data-gpu-memory-pressure="warning"] + .neuroglancer-layer-item-visible-progress, +.neuroglancer-layer-item[data-gpu-memory-pressure="warning"] + .neuroglancer-layer-item-prefetch-progress { + background-color: #d79b22; +} + +.neuroglancer-layer-item[data-gpu-memory-pressure="critical"] + .neuroglancer-layer-item-visible-progress, +.neuroglancer-layer-item[data-gpu-memory-pressure="critical"] + .neuroglancer-layer-item-prefetch-progress { + background-color: #e0523f; +} + .neuroglancer-layer-item-visible-progress { top: 0px; } diff --git a/src/ui/layer_bar.ts b/src/ui/layer_bar.ts index b4339954e8..b162b2be9b 100644 --- a/src/ui/layer_bar.ts +++ b/src/ui/layer_bar.ts @@ -17,11 +17,18 @@ import "#src/noselect.css"; import "#src/ui/layer_bar.css"; import svg_plus from "ikonate/icons/plus.svg?raw"; +import { throttle } from "lodash-es"; import type { ManagedUserLayer } from "#src/layer/index.js"; import { addNewLayer, deleteLayer, makeLayer } from "#src/layer/index.js"; import type { LayerGroupViewer } from "#src/layer_group_viewer.js"; import { NavigationLinkType } from "#src/navigation_state.js"; import type { WatchableValueInterface } from "#src/trackable_value.js"; +import { + computeGpuMemoryBytes, + getGpuMemoryPressure, + getLayerGpuMemoryPressure, + type GpuMemoryPressure, +} from "#src/ui/layer_bar_gpu_memory_pressure.js"; import type { DropLayers } from "#src/ui/layer_drag_and_drop.js"; import { registerLayerBarDragLeaveHandler, @@ -42,6 +49,8 @@ import { makeDeleteButton } from "#src/widget/delete_button.js"; import { makeIcon } from "#src/widget/icon.js"; import { PositionWidget } from "#src/widget/position_widget.js"; +const GPU_MEMORY_PRESSURE_UPDATE_INTERVAL_MS = 1000; + class LayerWidget extends RefCounted { element = document.createElement("div"); layerNumberElement = document.createElement("div"); @@ -252,6 +261,8 @@ export class LayerBar extends RefCounted { element = document.createElement("div"); private layerUpdateNeeded = true; private valueUpdateNeeded = false; + private gpuMemoryPressure: GpuMemoryPressure = "normal"; + private gpuMemoryPressureUpdatePending = false; dropZone: HTMLDivElement; private layerWidgetInsertionPoint = document.createElement("div"); private positionWidget: PositionWidget; @@ -373,6 +384,7 @@ export class LayerBar extends RefCounted { this.update(); this.updateChunkStatistics(); + this.scheduleGpuMemoryPressureUpdate(); registerLayerBarDragLeaveHandler(this); registerLayerBarDropHandlers(this, dropZone, undefined); @@ -386,10 +398,15 @@ export class LayerBar extends RefCounted { this.registerDisposer( manager.chunkManager.layerChunkStatisticsUpdated.add( this.registerCancellable( - animationFrameDebounce(() => this.updateChunkStatistics()), + animationFrameDebounce(() => this.handleChunkStatisticsUpdated()), ), ), ); + this.registerDisposer( + manager.chunkManager.chunkQueueManager.capacities.gpuMemory.sizeLimit.changed.add( + () => this.scheduleGpuMemoryPressureUpdate(), + ), + ); } disposed() { @@ -449,6 +466,46 @@ export class LayerBar extends RefCounted { } } + private handleChunkStatisticsUpdated() { + this.updateChunkStatistics(); + this.scheduleGpuMemoryPressureUpdate(); + } + + private scheduleGpuMemoryPressureUpdate = this.registerCancellable( + throttle( + () => { + void this.updateGpuMemoryPressure(); + }, + GPU_MEMORY_PRESSURE_UPDATE_INTERVAL_MS, + { leading: true, trailing: true }, + ), + ); + + private async updateGpuMemoryPressure() { + if (this.gpuMemoryPressureUpdatePending || this.wasDisposed) { + return; + } + this.gpuMemoryPressureUpdatePending = true; + let pressure: GpuMemoryPressure = "normal"; + try { + const chunkQueueManager = this.manager.chunkManager.chunkQueueManager; + const statistics = await chunkQueueManager.getStatistics(); + pressure = getGpuMemoryPressure( + computeGpuMemoryBytes(statistics.values()), + chunkQueueManager.capacities.gpuMemory.sizeLimit.value, + ); + } catch { + pressure = "normal"; + } finally { + this.gpuMemoryPressureUpdatePending = false; + } + if (this.wasDisposed || pressure === this.gpuMemoryPressure) { + return; + } + this.gpuMemoryPressure = pressure; + this.updateChunkStatistics(); + } + private updateChunkStatistics() { for (const [layer, widget] of this.layerWidgets) { let numVisibleChunksNeeded = 0; @@ -475,6 +532,11 @@ export class LayerBar extends RefCounted { (numPrefetchChunksAvailable / Math.max(1, numPrefetchChunksNeeded)) * 100 }%`; + widget.element.dataset.gpuMemoryPressure = getLayerGpuMemoryPressure( + this.gpuMemoryPressure, + numVisibleChunksNeeded, + numVisibleChunksAvailable, + ); } } diff --git a/src/ui/layer_bar_gpu_memory_pressure.spec.ts b/src/ui/layer_bar_gpu_memory_pressure.spec.ts new file mode 100644 index 0000000000..b4f2cdc2d8 --- /dev/null +++ b/src/ui/layer_bar_gpu_memory_pressure.spec.ts @@ -0,0 +1,67 @@ +import { describe, expect, it } from "vitest"; + +import { + ChunkMemoryStatistics, + ChunkPriorityTier, + ChunkState, + getChunkStateStatisticIndex, + numChunkMemoryStatistics, + numChunkStatistics, +} from "#src/chunk_manager/base.js"; +import { + computeGpuMemoryBytes, + getGpuMemoryPressure, + getLayerGpuMemoryPressure, +} from "#src/ui/layer_bar_gpu_memory_pressure.js"; + +function makeStatistics( + gpuMemoryBytesByTier: Partial>, +) { + const statistics = new Float64Array(numChunkStatistics); + for (const [tier, bytes] of Object.entries(gpuMemoryBytesByTier)) { + statistics[ + getChunkStateStatisticIndex(ChunkState.GPU_MEMORY, Number(tier)) * + numChunkMemoryStatistics + + ChunkMemoryStatistics.gpuMemoryBytes + ] = bytes; + } + return statistics; +} + +describe("ui/layer_bar_gpu_memory_pressure", () => { + it("sums GPU memory across all GPU_MEMORY priority tiers and sources", () => { + const firstSource = makeStatistics({ + [ChunkPriorityTier.VISIBLE]: 10, + [ChunkPriorityTier.PREFETCH]: 20, + [ChunkPriorityTier.RECENT]: 30, + }); + const secondSource = makeStatistics({ + [ChunkPriorityTier.VISIBLE]: 40, + [ChunkPriorityTier.RECENT]: 50, + }); + + expect(computeGpuMemoryBytes([firstSource, secondSource])).toBe(150); + }); + + it("maps GPU usage ratio to warning and critical pressure", () => { + expect(getGpuMemoryPressure(89, 100)).toBe("normal"); + expect(getGpuMemoryPressure(90, 100)).toBe("warning"); + expect(getGpuMemoryPressure(97, 100)).toBe("warning"); + expect(getGpuMemoryPressure(98, 100)).toBe("critical"); + }); + + it("treats invalid or unlimited GPU memory limits as normal pressure", () => { + expect(getGpuMemoryPressure(100, Number.POSITIVE_INFINITY)).toBe("normal"); + expect(getGpuMemoryPressure(100, Number.NaN)).toBe("normal"); + expect(getGpuMemoryPressure(100, 0)).toBe("normal"); + expect(getGpuMemoryPressure(Number.NaN, 100)).toBe("normal"); + }); + + it("only applies global pressure to layers missing visible chunks", () => { + expect(getLayerGpuMemoryPressure("warning", 10, 9)).toBe("warning"); + expect(getLayerGpuMemoryPressure("critical", 10, 9)).toBe("critical"); + expect(getLayerGpuMemoryPressure("warning", 10, 10)).toBe("normal"); + expect(getLayerGpuMemoryPressure("warning", 0, 0)).toBe("normal"); + expect(getLayerGpuMemoryPressure("normal", 10, 9)).toBe("normal"); + }); +}); diff --git a/src/ui/layer_bar_gpu_memory_pressure.ts b/src/ui/layer_bar_gpu_memory_pressure.ts new file mode 100644 index 0000000000..e010f93868 --- /dev/null +++ b/src/ui/layer_bar_gpu_memory_pressure.ts @@ -0,0 +1,86 @@ +/** + * @license + * Copyright 2026 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { + ChunkMemoryStatistics, + ChunkPriorityTier, + ChunkState, + getChunkStateStatisticIndex, + numChunkMemoryStatistics, +} from "#src/chunk_manager/base.js"; + +export type GpuMemoryPressure = "normal" | "warning" | "critical"; + +export const GPU_MEMORY_PRESSURE_WARNING_RATIO = 0.9; +export const GPU_MEMORY_PRESSURE_CRITICAL_RATIO = 0.98; + +export function computeGpuMemoryBytes( + chunkStatistics: Iterable, +) { + let total = 0; + for (const statistics of chunkStatistics) { + for ( + let tier = ChunkPriorityTier.FIRST_TIER; + tier <= ChunkPriorityTier.LAST_TIER; + ++tier + ) { + total += + statistics[ + getChunkStateStatisticIndex(ChunkState.GPU_MEMORY, tier) * + numChunkMemoryStatistics + + ChunkMemoryStatistics.gpuMemoryBytes + ] || 0; + } + } + return total; +} + +export function getGpuMemoryPressure( + gpuMemoryBytes: number, + gpuMemoryLimitBytes: number, +): GpuMemoryPressure { + if ( + !Number.isFinite(gpuMemoryBytes) || + !Number.isFinite(gpuMemoryLimitBytes) || + gpuMemoryBytes < 0 || + gpuMemoryLimitBytes <= 0 + ) { + return "normal"; + } + const ratio = gpuMemoryBytes / gpuMemoryLimitBytes; + if (ratio >= GPU_MEMORY_PRESSURE_CRITICAL_RATIO) { + return "critical"; + } + if (ratio >= GPU_MEMORY_PRESSURE_WARNING_RATIO) { + return "warning"; + } + return "normal"; +} + +export function getLayerGpuMemoryPressure( + globalPressure: GpuMemoryPressure, + numVisibleChunksNeeded: number, + numVisibleChunksAvailable: number, +): GpuMemoryPressure { + if ( + globalPressure === "normal" || + numVisibleChunksNeeded <= 0 || + numVisibleChunksAvailable >= numVisibleChunksNeeded + ) { + return "normal"; + } + return globalPressure; +}