Skip to content

Commit

Permalink
Perform GPU match only when we are going through the optimized profile (
Browse files Browse the repository at this point in the history
#107)

Signed-off-by: Shiva Krishna, Merla <[email protected]>
  • Loading branch information
shivamerla authored Aug 28, 2024
1 parent e43443b commit 785fbd7
Showing 1 changed file with 18 additions and 3 deletions.
21 changes: 18 additions & 3 deletions internal/nimparser/nimparser.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,11 @@ import (
"gopkg.in/yaml.v2"
)

const (
// BackendTypeTensorRT indicates tensortt backend
BackendTypeTensorRT = "tensorrt"
)

// File represents the model files
type File struct {
Name string `yaml:"name" json:"name,omitempty"`
Expand Down Expand Up @@ -152,11 +157,17 @@ func MatchProfiles(modelSpec appsv1alpha1.ModelSpec, manifest NIMManifest, disco
continue
}

// GPU matching logic
if len(modelSpec.GPUs) > 0 || len(discoveredGPUs) > 0 {
if !matchGPUProfile(modelSpec, profile, discoveredGPUs) {
// Perform GPU match only when optimized engine is selected or GPU filters are provided
if isOptimizedEngine(modelSpec.Engine) || len(modelSpec.GPUs) > 0 {
// Skip non optimized profiles
if !isOptimizedEngine(backend) {
continue
}
if len(modelSpec.GPUs) > 0 || len(discoveredGPUs) > 0 {
if !matchGPUProfile(modelSpec, profile, discoveredGPUs) {
continue
}
}
}

// Profile matched the given model parameters, add hash to the selected profiles
Expand All @@ -166,6 +177,10 @@ func MatchProfiles(modelSpec appsv1alpha1.ModelSpec, manifest NIMManifest, disco
return selectedProfiles, nil
}

func isOptimizedEngine(engine string) bool {
return engine != "" && strings.Contains(strings.ToLower(engine), BackendTypeTensorRT)
}

func matchGPUProfile(modelSpec appsv1alpha1.ModelSpec, profile NIMProfile, discoveredGPUs []string) bool {
foundGPU := false

Expand Down

0 comments on commit 785fbd7

Please sign in to comment.