Skip to content

feat(gpu): strict device affinity requirement and refactor #1121

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
JustinChengLZ wants to merge 2 commits into
base: main
Choose a base branch
from
Open

feat(gpu): strict device affinity requirement and refactor #1121

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9a39f17
feat: refactor device affinity and support strict device affinity req…
JustinChengLZ Apr 6, 2026
0d0a140
fix: simplify the device affinity logic
JustinChengLZ Apr 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions cmd/katalyst-agent/app/options/qrm/gpu_plugin.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ type GPUOptions struct {
GPUMemoryAllocatablePerGPU string
SkipGPUStateCorruption bool
RDMADeviceNames []string
RequiredDeviceAffinity bool

GPUStrategyOptions *gpustrategy.GPUStrategyOptions
}
Expand All @@ -41,6 +42,7 @@ func NewGPUOptions() *GPUOptions {
GPUMemoryAllocatablePerGPU: "100",
RDMADeviceNames: []string{},
GPUStrategyOptions: gpustrategy.NewGPUStrategyOptions(),
RequiredDeviceAffinity: true,
}
}

Expand All @@ -55,6 +57,8 @@ func (o *GPUOptions) AddFlags(fss *cliflag.NamedFlagSets) {
fs.BoolVar(&o.SkipGPUStateCorruption, "skip-gpu-state-corruption",
o.SkipGPUStateCorruption, "skip gpu state corruption, and it will be used after updating state properties")
fs.StringSliceVar(&o.RDMADeviceNames, "rdma-resource-names", o.RDMADeviceNames, "The name of the RDMA resource")
fs.BoolVar(&o.RequiredDeviceAffinity, "gpu-required-device-affinity", o.RequiredDeviceAffinity,
"required device affinity, and when true it will cause pods to admit fail if unable to meet device affinity")
o.GPUStrategyOptions.AddFlags(fss)
}

Expand All @@ -71,5 +75,6 @@ func (o *GPUOptions) ApplyTo(conf *qrmconfig.GPUQRMPluginConfig) error {
if err := o.GPUStrategyOptions.ApplyTo(conf.GPUStrategyConfig); err != nil {
return err
}
conf.RequiredDeviceAffinity = o.RequiredDeviceAffinity
return nil
}
2 changes: 1 addition & 1 deletion pkg/agent/qrm-plugins/gpu/baseplugin/base.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -246,7 +246,7 @@ func (p *BasePlugin) UpdateAllocatableAssociatedDevices(
deviceTopology.Devices[device.ID] = machine.DeviceInfo{
Health: device.Health,
NumaNodes: numaNode,
DeviceAffinity: make(map[machine.AffinityPriority]machine.DeviceIDs),
DeviceAffinity: make(map[machine.Dimension]machine.DeviceIDs),
}
}

Expand Down
Loading
Loading