Repository navigation
feat(rod): show per-user dominant resource allocation for selected queue - #1587
Closed
janzak-dev wants to merge 4 commits into
Closed
janzak-dev wants to merge 4 commits into
janzak-dev wants to merge 4 commits into
Conversation
Add a queue-wide breakdown of requested resources per experiment creator to GetRunningAndPendingJobs and render it on /resources as a stacked bar with one chip per user for the first resource picked in "I need". - storage: AggregateActiveResourcesByAuthor sums reque 8000 sted units by creator, canonical resource type and running/pending state across all active work units (Spanner and Postgres), and WithReadOnlyWorkUnitSnapshot binds it and FetchActiveWorkUnitsByQueues to one read-only snapshot. - api: GetRunningAndPendingJobsResponse.user_allocations (field 4). - rod_data_provider: populate user_allocations in the same snapshot as the jobs page; memory is normalized to GiB. - ui: QueueUserAllocationBar above the jobs table; generic parents such as nvidia.com/gpu sum all matching models via matchesResourceFilter. Bug: 582253355 TAG=agy CONV=4ed50232-a6f2-4c76-aa6f-ef1c3833a550
Remove the Repository-level AggregateActiveResourcesByAuthor (the RPC only reads through WithReadOnlyWorkUnitSnapshot), duplicate positive-unit filters and re-canonicalization, nil-callback guards, and the forwarding formatCompoundAccelerator wrapper; restore dropped comments. Bug: 582253355 TAG=agy CONV=4ed50232-a6f2-4c76-aa6f-ef1c3833a550
Bug: 582253355 TAG=agy CONV=4ed50232-a6f2-4c76-aa6f-ef1c3833a550
…uards Unexport package-internal canonicalization and state classification helpers, drop context and nil checks already enforced by callers, remove an unused mock field and summary property, and drop a zero-total filter for allocations the backend never emits. Bug: 582253355 TAG=agy CONV=4ed50232-a6f2-4c76-aa6f-ef1c3833a550
There was a problem hiding this comment.
Code Review
This pull request introduces a new feature to aggregate and display per-user resource allocations (running and pending) for jobs in a queue. It adds a new UserResourceAllocation proto message, updates the rod_data_provider to calculate these allocations using a consistent storage snapshot, and adds a new UI component to visualize the data. My feedback suggests adding a nil check for the cluster queue retrieval to improve robustness.
Comment on lines
+42
to
+45
| cq, ok := kueueCache.ClusterQueues.Get(clusterQueueName) | ||
| if !ok { | ||
| return nil | ||
| } |
There was a problem hiding this comment.
To prevent a potential nil pointer dereference panic, add a nil check for cq in case the cache retrieves a nil entry for the cluster queue name.
Suggested change
| cq, ok := kueueCache.ClusterQueues.Get(clusterQueueName) | |
| if !ok { | |
| return nil | |
| } | |
| cq, ok := kueueCache.ClusterQueues.Get(clusterQueueName) | |
| if !ok || cq == nil { | |
| return nil | |
| } |
Collaborator
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The Resource Observability Dashboard (
/resources) now shows how much of the dominant resource each user holds in the selected queue. The dominant resource is the first resource picked in the Smart Resource Finder's "I need" picker. A stacked bar with one chip per user sits above the jobs table, so you can see who to ping without paging through jobs.Motivation & Context
Part of #1515 (the "Who holds the used GPUs?" bar in the mockups). The jobs table is paginated, so answering "who is using this queue?" meant clicking through every page and adding up the Resources column by hand. The breakdown is computed server-side across all active work units of the queue, so it doesn't depend on pagination. It is not hardcoded to GPUs: the dominant resource can be a specific accelerator model (
nvidia.com/gpu: nvidia-tesla-h100), a generic parent (nvidia.com/gpu, which sums every matching model), orcpu/memory.Key Changes
UserResourceAllocation { user, resource_type, running_units, pending_units }andGetRunningAndPendingJobsResponse.user_allocations = 4. All resource types are returned and the UI picks what to show.WorkUnitSnapshotRepositorywithFetchActiveWorkUnitsByQueuesandAggregateActiveResourcesByAuthor, plusWorkUnitStore.WithReadOnlyWorkUnitSnapshot, so the jobs page and the aggregation are read from one read-only snapshot (SpannerReadOnlyTransaction, PostgresREPEATABLE READ READ ONLY). Rows are grouped byExperiments.Creator, canonical resource type and running vs pending. The existingFetchActiveWorkUnitsByQueuesnow delegates to the snapshot. Shared aggregation lives ininternal/storage/data;MockStorageis updated.internal/common/resourceaccumulator/canonical.gomaps symbolic GPU types, GKE accelerator labels and TPU models to the same compound names Kueue queues use (nvidia.com/gpu: <model>,google.com/tpu: <model>) and normalizes memory to GiB.FormatCompoundAcceleratorandBytesPerGiBmove here fromrod_data_provider, removing a duplicate.rod_data_provider:GetRunningAndPendingJobsfillsuser_allocationsin the same snapshot as the jobs page, sorted by total units. A genericnvidia.com/gpurequest is attributed to the queue's GPU model when the queue has exactly one GPU flavor.QueueUserAllocationBar(signals,OnPush,@if/@for, harness tests) rendered inside the jobs card above the table. It shows a stacked bar plus display-only chips with running and pending units and the unit label (GPUs / chips / cores / GiB), and filters with the existingmatchesResourceFilter. With no resource picked it shows a hint; when nothing matches it shows an empty message.ResourceDashboardpassesdominantResourcetoQueueJobsTable.docs/flows/metrics_and_observability.mdsection 5.Assumptions
RUNNING/STARTING(admitted by Kueue, holding quota); pending =SCHEDULING/INITIALIZING. This follows the reconciler's state mapping and the status-history "admitted" states. If wrong, units shift between the running and pending labels; per-user totals are unaffected.GetRunningAndPendingJobs; a follow-up could cap it or precompute it.canonical.gomirrors the table used when launching on GKE. If that table changes without this one, those requests show up under a different model name, though they still sum under the genericnvidia.com/gpuparent.nvidia.com/gpu. If wrong, a specific-model filter can miss generic requests in single-flavor queues (the parent filter always includes them).unknownin the UI.WorkUnitStoreinterface forces both backends into the same change. A vertical split was evaluated (API+storage → UI → snapshot/flavor resolution) and can still be applied if reviewers prefer.Verification & Testing
bazel teston//internal/common/resourceaccumulator/...,//internal/storage/...(Spanner emulator and Postgres),//cmd/rod_data_provider/...and//cmd/ui:ui_unit_tests;./.githooks/pre-commit(formatters, linters,run_affected_tests.py: 300/300 affected tests pass);npx playwright test screenshot_tests/resource_dashboard(existing/resourcesbaselines pass).cpu/memory(GiB) / GPU (symbolic and GKE-label types) / TPU, and 6 active work units vs a 2-row first page inside the same snapshot. This shows the aggregation is independent of pagination.user_allocationsmapping, ordering and tie-breakers, single-flavor and multi-flavor GPU attribution, and one snapshot per request.cpu(cores) andmemory(GiB), empty-picker hint, no-match message, loading state, and a dashboard test showing that changing the picker updates the bar.active_resources.go100%,user_allocations.go~100%,jobs.go89%,canonical.go~95%, Postgres / Spanner aggregation 82% / 76%. The uncovered lines are database scan and unmarshal error-wrapping paths.user_allocationsequals the per-user sum of requested units over all pages of jobs. For example, a CPU queue with 14 active jobs (4 × 0.5 cores + 10 × 100 cores, 0.5 GiB each, allSCHEDULING/INITIALIZING) returnedcpu 0 running + 1002 pendingandmemory 0 + 7 GiB. Withpage_size=5, all three pages returned the same queue-wide totals.cpu: 0.5, 10 atcpu: 100).Related Issues / Bugs
Part of #1515
TAG=agy
CONV=4ed50232-a6f2-4c76-aa6f-ef1c3833a550