|
| 1 | +/* |
| 2 | + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
| 3 | + * SPDX-License-Identifier: Apache-2.0 |
| 4 | + * |
| 5 | + * |
| 6 | + * Licensed under the Apache License, Version 2.0 (the "License"); |
| 7 | + * you may not use this file except in compliance with the License. |
| 8 | + * You may obtain a copy of the License at |
| 9 | + * |
| 10 | + * http://www.apache.org/licenses/LICENSE-2.0 |
| 11 | + * |
| 12 | + * Unless required by applicable law or agreed to in writing, software |
| 13 | + * distributed under the License is distributed on an "AS IS" BASIS, |
| 14 | + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 15 | + * See the License for the specific language governing permissions and |
| 16 | + * limitations under the License. |
| 17 | + */ |
| 18 | + |
| 19 | +package app |
| 20 | + |
| 21 | +import ( |
| 22 | + "fmt" |
| 23 | + |
| 24 | + "github.com/spf13/cobra" |
| 25 | + |
| 26 | + "github.com/NVIDIA/skyhook/operator/internal/cli/client" |
| 27 | + cliContext "github.com/NVIDIA/skyhook/operator/internal/cli/context" |
| 28 | + "github.com/NVIDIA/skyhook/operator/internal/cli/utils" |
| 29 | +) |
| 30 | + |
| 31 | +// lifecycleConfig defines the configuration for a lifecycle command |
| 32 | +type lifecycleConfig struct { |
| 33 | + use string |
| 34 | + short string |
| 35 | + long string |
| 36 | + example string |
| 37 | + annotation string |
| 38 | + action string // "set" or "remove" |
| 39 | + verb string // past tense for output message (e.g., "paused", "resumed") |
| 40 | + confirmVerb string // verb for confirmation prompt (e.g., "pause", "disable") |
| 41 | + needsConfirm bool |
| 42 | +} |
| 43 | + |
| 44 | +// lifecycleOptions holds the options for lifecycle commands that need confirmation |
| 45 | +type lifecycleOptions struct { |
| 46 | + confirm bool |
| 47 | +} |
| 48 | + |
| 49 | +// newLifecycleCmd creates a lifecycle command based on the provided configuration |
| 50 | +func newLifecycleCmd(ctx *cliContext.CLIContext, cfg lifecycleConfig) *cobra.Command { |
| 51 | + opts := &lifecycleOptions{} |
| 52 | + |
| 53 | + cmd := &cobra.Command{ |
| 54 | + Use: cfg.use, |
| 55 | + Short: cfg.short, |
| 56 | + Long: cfg.long, |
| 57 | + Example: cfg.example, |
| 58 | + Args: cobra.ExactArgs(1), |
| 59 | + RunE: func(cmd *cobra.Command, args []string) error { |
| 60 | + skyhookName := args[0] |
| 61 | + |
| 62 | + if cfg.needsConfirm && !opts.confirm { |
| 63 | + _, _ = fmt.Fprintf(cmd.OutOrStdout(), "This will %s Skyhook %q. Continue? [y/N]: ", |
| 64 | + cfg.confirmVerb, skyhookName) |
| 65 | + var response string |
| 66 | + if _, err := fmt.Scanln(&response); err != nil || (response != "y" && response != "Y") { |
| 67 | + _, _ = fmt.Fprintln(cmd.OutOrStdout(), "Aborted.") |
| 68 | + return nil |
| 69 | + } |
| 70 | + } |
| 71 | + |
| 72 | + clientFactory := client.NewFactory(ctx.GlobalFlags.ConfigFlags) |
| 73 | + kubeClient, err := clientFactory.Client() |
| 74 | + if err != nil { |
| 75 | + return fmt.Errorf("initializing kubernetes client: %w", err) |
| 76 | + } |
| 77 | + |
| 78 | + if cfg.action == "set" { |
| 79 | + err = utils.SetSkyhookAnnotation(cmd.Context(), kubeClient.Dynamic(), skyhookName, cfg.annotation, "true") |
| 80 | + } else { |
| 81 | + err = utils.RemoveSkyhookAnnotation(cmd.Context(), kubeClient.Dynamic(), skyhookName, cfg.annotation) |
| 82 | + } |
| 83 | + if err != nil { |
| 84 | + return err |
| 85 | + } |
| 86 | + |
| 87 | + _, _ = fmt.Fprintf(cmd.OutOrStdout(), "Skyhook %q %s\n", skyhookName, cfg.verb) |
| 88 | + return nil |
| 89 | + }, |
| 90 | + } |
| 91 | + |
| 92 | + if cfg.needsConfirm { |
| 93 | + cmd.Flags().BoolVarP(&opts.confirm, "confirm", "y", false, "Skip confirmation prompt") |
| 94 | + } |
| 95 | + |
| 96 | + return cmd |
| 97 | +} |
| 98 | + |
| 99 | +// NewPauseCmd creates the pause command |
| 100 | +func NewPauseCmd(ctx *cliContext.CLIContext) *cobra.Command { |
| 101 | + return newLifecycleCmd(ctx, lifecycleConfig{ |
| 102 | + use: "pause <skyhook-name>", |
| 103 | + short: "Pause a Skyhook from processing", |
| 104 | + long: `Pause a Skyhook by setting the pause annotation. |
| 105 | +
|
| 106 | +When a Skyhook is paused, the operator will stop processing new nodes |
| 107 | +but will not interrupt any currently running operations.`, |
| 108 | + example: ` # Pause a Skyhook |
| 109 | + kubectl skyhook pause gpu-init |
| 110 | +
|
| 111 | + # Pause without confirmation |
| 112 | + kubectl skyhook pause gpu-init --confirm`, |
| 113 | + annotation: utils.PauseAnnotation, |
| 114 | + action: "set", |
| 115 | + verb: "paused", |
| 116 | + confirmVerb: "pause", |
| 117 | + needsConfirm: true, |
| 118 | + }) |
| 119 | +} |
| 120 | + |
| 121 | +// NewResumeCmd creates the resume command |
| 122 | +func NewResumeCmd(ctx *cliContext.CLIContext) *cobra.Command { |
| 123 | + return newLifecycleCmd(ctx, lifecycleConfig{ |
| 124 | + use: "resume <skyhook-name>", |
| 125 | + short: "Resume a paused Skyhook", |
| 126 | + long: `Resume a paused Skyhook by removing the pause annotation. |
| 127 | +
|
| 128 | +The operator will resume processing nodes after this command.`, |
| 129 | + example: ` # Resume a paused Skyhook |
| 130 | + kubectl skyhook resume gpu-init`, |
| 131 | + annotation: utils.PauseAnnotation, |
| 132 | + action: "remove", |
| 133 | + verb: "resumed", |
| 134 | + needsConfirm: false, |
| 135 | + }) |
| 136 | +} |
| 137 | + |
| 138 | +// NewDisableCmd creates the disable command |
| 139 | +func NewDisableCmd(ctx *cliContext.CLIContext) *cobra.Command { |
| 140 | + return newLifecycleCmd(ctx, lifecycleConfig{ |
| 141 | + use: "disable <skyhook-name>", |
| 142 | + short: "Disable a Skyhook completely", |
| 143 | + long: `Disable a Skyhook by setting the disable annotation. |
| 144 | +
|
| 145 | +When a Skyhook is disabled, the operator will completely stop processing |
| 146 | +and the Skyhook will be effectively inactive.`, |
| 147 | + example: ` # Disable a Skyhook |
| 148 | + kubectl skyhook disable gpu-init |
| 149 | +
|
| 150 | + # Disable without confirmation |
| 151 | + kubectl skyhook disable gpu-init --confirm`, |
| 152 | + annotation: utils.DisableAnnotation, |
| 153 | + action: "set", |
| 154 | + verb: "disabled", |
| 155 | + confirmVerb: "disable", |
| 156 | + needsConfirm: true, |
| 157 | + }) |
| 158 | +} |
| 159 | + |
| 160 | +// NewEnableCmd creates the enable command |
| 161 | +func NewEnableCmd(ctx *cliContext.CLIContext) *cobra.Command { |
| 162 | + return newLifecycleCmd(ctx, lifecycleConfig{ |
| 163 | + use: "enable <skyhook-name>", |
| 164 | + short: "Enable a disabled Skyhook", |
| 165 | + long: `Enable a disabled Skyhook by removing the disable annotation. |
| 166 | +
|
| 167 | +The operator will resume normal processing after this command.`, |
| 168 | + example: ` # Enable a disabled Skyhook |
| 169 | + kubectl skyhook enable gpu-init`, |
| 170 | + annotation: utils.DisableAnnotation, |
| 171 | + action: "remove", |
| 172 | + verb: "enabled", |
| 173 | + needsConfirm: false, |
| 174 | + }) |
| 175 | +} |
0 commit comments