diff --git a/kwok/README.md b/kwok/README.md index a3166d0bb8..bf168c0a44 100644 --- a/kwok/README.md +++ b/kwok/README.md @@ -71,9 +71,8 @@ After doing this, you can create a deployment to test node scaling with kwok pro ## Specifying Instance Types By default, the KWOK provider will create a hypothetical set of instance types that it uses for node provisioning. You -can specify a custom set of instance types by providing a JSON file with the list of supported instance options. This -set of instance types is embedded into the binary on creation; if you want to change the instance types that -Karpenter+KWOK support, you will need to adjust the embedded data and recompile. +can specify a custom set of instance types by providing a JSON file with the list of supported instance options. To do so, +set the `--instance-types-file-path` flag or `INSTANCE_TYPES_FILE_PATH` environment variable to your custom file's path. There is an example instance types file in [examples/instance\_types.json](examples/instance_types.json) that you can regenerate with `make gen_instance_types`. diff --git a/kwok/cloudprovider/helpers.go b/kwok/cloudprovider/helpers.go index 680e097bde..b93430854b 100644 --- a/kwok/cloudprovider/helpers.go +++ b/kwok/cloudprovider/helpers.go @@ -20,6 +20,7 @@ import ( _ "embed" "encoding/json" "fmt" + "os" "regexp" "github.com/samber/lo" @@ -61,10 +62,29 @@ type InstanceTypeOptions struct { } //go:embed instance_types.json -var rawInstanceTypes []byte +var defaultRawInstanceTypes []byte + +type InstanceTypesOptions struct { + CustomInstanceTypesFilePath *string +} + +type InstanceTypesOption func(*InstanceTypesOptions) // ConstructInstanceTypes create many instance types based on the embedded instance type data -func ConstructInstanceTypes() ([]*cloudprovider.InstanceType, error) { +func ConstructInstanceTypes(opts ...InstanceTypesOption) ([]*cloudprovider.InstanceType, error) { + o := &InstanceTypesOptions{} + for _, opt := range opts { + opt(o) + } + rawInstanceTypes := defaultRawInstanceTypes + if o.CustomInstanceTypesFilePath != nil { + customRawInstanceTypes, err := os.ReadFile(*o.CustomInstanceTypesFilePath) + if err != nil { + return nil, fmt.Errorf("could not read custom instance types file: %w", err) + } + rawInstanceTypes = customRawInstanceTypes + } + var instanceTypes []*cloudprovider.InstanceType var instanceTypeOptions []InstanceTypeOptions @@ -79,6 +99,13 @@ func ConstructInstanceTypes() ([]*cloudprovider.InstanceType, error) { return instanceTypes, nil } +// WithInstanceTypesFromFile constructs instance types from a custom file. +func WithInstanceTypesFromFile(path string) InstanceTypesOption { + return func(o *InstanceTypesOptions) { + o.CustomInstanceTypesFilePath = &path + } +} + // parseSizeFromType will attempt to discover the instance size if it matches a special AWS format. // If not, fall back to the cpu value. This works for both Azure and GCP (and the "generic" instances // generated by tools/gen_instances.go) diff --git a/kwok/main.go b/kwok/main.go index a84c6e0a0f..f09c332c01 100644 --- a/kwok/main.go +++ b/kwok/main.go @@ -26,7 +26,14 @@ import ( func main() { ctx, op := operator.NewOperator() - instanceTypes, err := kwok.ConstructInstanceTypes() + + options := kwok.GetOptionsOrDie() + + var instanceTypesOptions []kwok.InstanceTypesOption + if options.InstanceTypesFilePath != "" { + instanceTypesOptions = append(instanceTypesOptions, kwok.WithInstanceTypesFromFile(options.InstanceTypesFilePath)) + } + instanceTypes, err := kwok.ConstructInstanceTypes(instanceTypesOptions...) if err != nil { log.FromContext(ctx).Error(err, "failed constructing instance types") } diff --git a/kwok/options/options.go b/kwok/options/options.go new file mode 100644 index 0000000000..1799b61e5f --- /dev/null +++ b/kwok/options/options.go @@ -0,0 +1,69 @@ +/* +Copyright The Kubernetes Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package options + +import ( + "context" + "errors" + "flag" + "fmt" + "os" + + "sigs.k8s.io/karpenter/pkg/operator/options" + "sigs.k8s.io/karpenter/pkg/utils/env" +) + +func init() { + options.Injectables = append(options.Injectables, &Options{}) +} + +type optionsKey struct{} + +// Options contains all CLI flags / env vars for the KWOK cloudprovider. +type Options struct { + InstanceTypesFilePath string +} + +func (o *Options) AddFlags(fs *options.FlagSet) { + fs.StringVar(&o.InstanceTypesFilePath, "instance-types-file-path", env.WithDefaultString("INSTANCE_TYPES_FILE_PATH", ""), "Path to a custom instance-types file") +} + +func (o *Options) Parse(fs *options.FlagSet, args ...string) error { + if err := fs.Parse(args); err != nil { + if errors.Is(err, flag.ErrHelp) { + os.Exit(0) + } + return fmt.Errorf("parsing flags, %w", err) + } + return nil +} + +func (o *Options) ToContext(ctx context.Context) context.Context { + return ToContext(ctx, o) +} + +func ToContext(ctx context.Context, opts *Options) context.Context { + return context.WithValue(ctx, optionsKey{}, opts) +} + +func FromContext(ctx context.Context) *Options { + retval := ctx.Value(optionsKey{}) + if retval == nil { + return nil + } + return retval.(*Options) +}