Add --no-cleanup-on-exit option to GFD

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]>
NVIDIA · Sep 17, 2024 · 08363f4 · 08363f4
1 parent 71c1fa7
commit 08363f4
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 0 deletions.
diff --git a/api/config/v1/flags.go b/api/config/v1/flags.go
@@ -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"`
@@ -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)
 			}
 		}
 	}

diff --git a/api/config/v1/flags_test.go b/api/config/v1/flags_test.go
@@ -184,6 +184,7 @@ func TestMarshalFlags(t *testing.T) {
 				"gfd": {
 					"oneshot": null,
 					"noTimestamp": null,
+					"noCleanupOnExit": null,
 					"outputFile": null,
 					"sleepInterval": "0s",
 					"machineTypeFile": null
@@ -208,6 +209,7 @@ func TestMarshalFlags(t *testing.T) {
 				"gfd": {
 					"oneshot": null,
 					"noTimestamp": null,
+					"noCleanupOnExit": null,
 					"outputFile": null,
 					"sleepInterval": "5ns",
 					"machineTypeFile": null

diff --git a/cmd/gpu-feature-discovery/main.go b/cmd/gpu-feature-discovery/main.go
@@ -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,
@@ -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
 		}