Skip to content

Commit

Permalink
Add --no-cleanup-on-exit option to GFD
Browse files Browse the repository at this point in the history
This change adds a --no-cleanup-on-exit option to GFD that skips the
removal of the generated labels file.

Signed-off-by: Evan Lezar <[email protected]>
  • Loading branch information
elezar committed Sep 17, 2024
1 parent 71c1fa7 commit 08363f4
Show file tree
Hide file tree
Showing 3 changed files with 14 additions and 0 deletions.
3 changes: 3 additions & 0 deletions api/config/v1/flags.go
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,7 @@ func (f *deviceListStrategyFlag) UnmarshalJSON(b []byte) error {
type GFDCommandLineFlags struct {
Oneshot *bool `json:"oneshot" yaml:"oneshot"`
NoTimestamp *bool `json:"noTimestamp" yaml:"noTimestamp"`
NoCleanupOnExit *bool `json:"noCleanupOnExit" yaml:"noCleanupOnExit"`
SleepInterval *Duration `json:"sleepInterval" yaml:"sleepInterval"`
OutputFile *string `json:"outputFile" yaml:"outputFile"`
MachineTypeFile *string `json:"machineTypeFile" yaml:"machineTypeFile"`
Expand Down Expand Up @@ -168,6 +169,8 @@ func (f *Flags) UpdateFromCLIFlags(c *cli.Context, flags []cli.Flag) {
updateFromCLIFlag(&f.GFD.NoTimestamp, c, n)
case "machine-type-file":
updateFromCLIFlag(&f.GFD.MachineTypeFile, c, n)
case "no-cleanup-on-exit":
updateFromCLIFlag(&f.GFD.NoCleanupOnExit, c, n)
}
}
}
Expand Down
2 changes: 2 additions & 0 deletions api/config/v1/flags_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -184,6 +184,7 @@ func TestMarshalFlags(t *testing.T) {
"gfd": {
"oneshot": null,
"noTimestamp": null,
"noCleanupOnExit": null,
"outputFile": null,
"sleepInterval": "0s",
"machineTypeFile": null
Expand All @@ -208,6 +209,7 @@ func TestMarshalFlags(t *testing.T) {
"gfd": {
"oneshot": null,
"noTimestamp": null,
"noCleanupOnExit": null,
"outputFile": null,
"sleepInterval": "5ns",
"machineTypeFile": null
Expand Down
9 changes: 9 additions & 0 deletions cmd/gpu-feature-discovery/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,12 @@ func main() {
Usage: "Do not add the timestamp to the labels",
EnvVars: []string{"GFD_NO_TIMESTAMP"},
},
&cli.BoolFlag{
Name: "no-cleanup-on-exit",
Value: false,
Usage: "Do not remove the generated label file",
EnvVars: []string{"GFD_NO_CLEANUP_ON_EXIT"},
},
&cli.DurationFlag{
Name: "sleep-interval",
Value: 60 * time.Second,
Expand Down Expand Up @@ -229,6 +235,9 @@ type gfd struct {

func (d *gfd) run(sigs chan os.Signal) (bool, error) {
defer func() {
if d.config.Flags.GFD.NoCleanupOnExit != nil && *d.config.Flags.GFD.NoCleanupOnExit {
return
}
if d.config.Flags.UseNodeFeatureAPI != nil && *d.config.Flags.UseNodeFeatureAPI {
return
}
Expand Down

0 comments on commit 08363f4

Please sign in to comment.