[harishxr]

Summary

Adds per-GPU memory accounting to the agent's HostResourceManager so GPU memory can be tracked and admitted per physical GPU, as the local scheduling foundation for NVIDIA MPS GPU sharing. Each GPU is modeled as a small account {total, remaining, mode} (mode ∈ FREE / MPS / WHOLE_GPU). MPS tasks reserve a memory share on a GPU while a whole-GPU task takes it exclusively. This replaces the GPU allocation role previously served by the GPU StringSet.

Implementation details

• New agent/engine/host_resource_manager_gpu.go: a per-UUID pool (map[uuid]*gpuDeviceState) with {total, remaining, mode}. Admission keys off mode, which makes MPS/whole-GPU mode-mixing on one GPU impossible by construction. A GPU with no reported memory is seeded total 0 and is only usable via the whole-GPU path.
• Task GPU memory demand is a dedicated task.GPUMemoryDemand{MiB, WholeGPU} type, returned from ToHostResources() as a second map alongside the CPU/MEMORY/PORTS resource map (not encoded into the SDK ecstypes.Resource). consume / consumable / release / checkResourcesHealth take that demand map. A non-MPS container sets WholeGPU; MPS containers on a GPU sum their caps into MiB. There is no per-task stored balance, consume and release both recompute from the task, so they stay symmetric and restart-safe.
• doStart seeds one pool entry per GPU from the RCI-reported per-device memory, as an INTEGER capacity keyed GPU_MEMORY:<uuid>. A GPU whose memory is not reported (older ecs-init, or an NVML query failure) is seeded with capacity 0: whole-GPU tasks still run against it, MPS sharing is refused (a cap cannot be enforced without a known total). This preserves the pre-existing whole-GPU behavior, which never depended on memory being known.
• Removes the GPU StringSet's allocation/validation role now that the pool is the source of truth for GPU admission (the StringSet is no longer used for accounting).

Testing

• Built a custom agent from this change and installed it on a AL2023 GPU AMI instance, then ran tasks through the accounting paths.
• Watching the agent's per-GPU pool log confirmed the expected behavior. All GPUs seeded independently at their reported memory. A whole-GPU task drove its GPU to remaining 0 / WHOLE_GPU while the others stayed untouched.
• A second whole-GPU task on the same GPU was refused and stayed pending until the first stopped, after which the GPU returned to FREE and the queued task consumed it.
• Two MPS containers pinned to one GPU consumed a single summed reservation (remaining equals total minus the sum of caps).
• An MPS-shared GPU and a whole-GPU on a different device ran simultaneously.
• Mode isolation held both ways, so an MPS task was refused on a whole-GPU device and a whole-GPU task was refused on an MPS device.
• After an agent restart the pool rebuilt to the same per-GPU state by re-consuming the running tasks, with no negative accounting.

New tests cover the changes: yes

Description for the changelog

Enhancement - Add per-GPU memory accounting to host resource manager

Additional Information

Does this PR include breaking model changes? If so, Have you added transformation functions? No

Does this PR include the addition of new environment variables in the README? No

Licensing

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.