diff --git a/cmd/scheduler.go b/cmd/scheduler.go new file mode 100644 index 000000000..32fd48cc8 --- /dev/null +++ b/cmd/scheduler.go @@ -0,0 +1,39 @@ +// +// DISCLAIMER +// +// Copyright 2024 ArangoDB GmbH, Cologne, Germany +// +// 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. +// +// Copyright holder is ArangoDB GmbH, Cologne, Germany +// + +package cmd + +import ( + "github.com/spf13/cobra" + + "github.com/arangodb/kube-arangodb/pkg/scheduler" +) + +func init() { + cmd := &cobra.Command{ + Use: "scheduler", + } + + if err := scheduler.InitCommand(cmd); err != nil { + panic(err.Error()) + } + + cmdMain.AddCommand(cmd) +} diff --git a/pkg/scheduler/cli.go b/pkg/scheduler/cli.go new file mode 100644 index 000000000..3a9c31a2f --- /dev/null +++ b/pkg/scheduler/cli.go @@ -0,0 +1,150 @@ +// +// DISCLAIMER +// +// Copyright 2024 ArangoDB GmbH, Cologne, Germany +// +// 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. +// +// Copyright holder is ArangoDB GmbH, Cologne, Germany +// + +package scheduler + +import ( + "context" + "os" + "strings" + + "github.com/spf13/cobra" + "sigs.k8s.io/yaml" + + "github.com/arangodb/kube-arangodb/pkg/logging" + "github.com/arangodb/kube-arangodb/pkg/util" + "github.com/arangodb/kube-arangodb/pkg/util/constants" + "github.com/arangodb/kube-arangodb/pkg/util/errors" + "github.com/arangodb/kube-arangodb/pkg/util/kclient" +) + +func InitCommand(cmd *cobra.Command) error { + var c cli + return c.register(cmd) +} + +type cli struct { + Namespace string + + Labels []string + Envs []string + + Profiles []string + + Container string + + Image string +} + +func (c *cli) asRequest(args ...string) (Request, error) { + var r = Request{ + Labels: map[string]string{}, + Envs: map[string]string{}, + } + + for _, l := range c.Labels { + p := strings.SplitN(l, "=", 2) + if len(p) == 1 { + r.Labels[p[0]] = "" + logger.Debug("Label Discovered: %s", p[0]) + } else { + r.Labels[p[0]] = p[1] + logger.Debug("Label Discovered: %s=%s", p[0], p[1]) + } + } + + for _, l := range c.Envs { + p := strings.SplitN(l, "=", 2) + if len(p) == 1 { + return r, errors.Errorf("Missing value for env: %s", p[0]) + } else { + r.Envs[p[0]] = p[1] + logger.Debug("Env Discovered: %s=%s", p[0], p[1]) + } + } + + if len(c.Profiles) > 0 { + r.Profiles = c.Profiles + logger.Debug("Enabling profiles: %s", strings.Join(c.Profiles, ", ")) + } + + r.Container = util.NewType(c.Container) + if c.Image != "" { + r.Image = util.NewType(c.Image) + } + + r.Args = args + + return r, nil +} + +func (c *cli) register(cmd *cobra.Command) error { + if err := logging.Init(cmd); err != nil { + return err + } + + cmd.RunE = c.run + + f := cmd.PersistentFlags() + + f.StringVarP(&c.Namespace, "namespace", "n", constants.NamespaceWithDefault("default"), "Kubernetes namespace") + f.StringSliceVarP(&c.Labels, "label", "l", nil, "Scheduler Render Labels in format =") + f.StringSliceVarP(&c.Envs, "env", "e", nil, "Scheduler Render Envs in format =") + f.StringSliceVarP(&c.Profiles, "profile", "p", nil, "Scheduler Render Profiles") + f.StringVar(&c.Container, "container", DefaultContainerName, "Container Name") + f.StringVar(&c.Image, "image", "", "Image") + + return nil +} + +func (c *cli) run(cmd *cobra.Command, args []string) error { + if err := logging.Enable(); err != nil { + return err + } + + r, err := c.asRequest() + if err != nil { + return err + } + + k, ok := kclient.GetDefaultFactory().Client() + if !ok { + return errors.Errorf("Unable to create Kubernetes Client") + } + + s := NewScheduler(k, c.Namespace) + + rendered, profiles, err := s.Render(context.Background(), r) + if err != nil { + return err + } + logger.Debug("Enabled profiles: %s", strings.Join(profiles, ", ")) + + data, err := yaml.Marshal(rendered) + if err != nil { + return err + } + + if _, err := util.WriteAll(os.Stdout, data); err != nil { + return err + } + + return nil +} diff --git a/pkg/scheduler/input.go b/pkg/scheduler/input.go new file mode 100644 index 000000000..03f3c885c --- /dev/null +++ b/pkg/scheduler/input.go @@ -0,0 +1,91 @@ +// +// DISCLAIMER +// +// Copyright 2024 ArangoDB GmbH, Cologne, Germany +// +// 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. +// +// Copyright holder is ArangoDB GmbH, Cologne, Germany +// + +package scheduler + +import ( + "math" + + core "k8s.io/api/core/v1" + + schedulerApi "github.com/arangodb/kube-arangodb/pkg/apis/scheduler/v1alpha1" + schedulerContainerApi "github.com/arangodb/kube-arangodb/pkg/apis/scheduler/v1alpha1/container" + schedulerContainerResourcesApi "github.com/arangodb/kube-arangodb/pkg/apis/scheduler/v1alpha1/container/resources" + schedulerPodApi "github.com/arangodb/kube-arangodb/pkg/apis/scheduler/v1alpha1/pod" + schedulerPodResourcesApi "github.com/arangodb/kube-arangodb/pkg/apis/scheduler/v1alpha1/pod/resources" + "github.com/arangodb/kube-arangodb/pkg/util" +) + +const DefaultContainerName = "job" + +type Request struct { + Labels map[string]string + + Profiles []string + + Envs map[string]string + + Container *string + + Image *string + + Args []string +} + +func (r Request) AsTemplate() *schedulerApi.ProfileTemplate { + var container schedulerContainerApi.Container + + if len(r.Envs) > 0 { + container.Environments = &schedulerContainerResourcesApi.Environments{} + + for k, v := range r.Envs { + container.Environments.Env = append(container.Environments.Env, core.EnvVar{ + Name: k, + Value: v, + }) + } + } + + if len(r.Args) > 0 { + container.Core = &schedulerContainerResourcesApi.Core{ + Args: r.Args, + } + } + + if r.Image != nil { + container.Image = &schedulerContainerResourcesApi.Image{ + Image: util.NewType(util.TypeOrDefault(r.Image)), + } + } + + return &schedulerApi.ProfileTemplate{ + Priority: util.NewType(math.MaxInt), + Pod: &schedulerPodApi.Pod{ + Metadata: &schedulerPodResourcesApi.Metadata{ + Labels: util.MergeMaps(true, r.Labels), + }, + }, + Container: &schedulerApi.ProfileContainerTemplate{ + Containers: map[string]schedulerContainerApi.Container{ + util.TypeOrDefault(r.Container, DefaultContainerName): container, + }, + }, + } +} diff --git a/pkg/scheduler/logger.go b/pkg/scheduler/logger.go new file mode 100644 index 000000000..b470e9a2a --- /dev/null +++ b/pkg/scheduler/logger.go @@ -0,0 +1,25 @@ +// +// DISCLAIMER +// +// Copyright 2024 ArangoDB GmbH, Cologne, Germany +// +// 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. +// +// Copyright holder is ArangoDB GmbH, Cologne, Germany +// + +package scheduler + +import "github.com/arangodb/kube-arangodb/pkg/logging" + +var logger = logging.Global().RegisterAndGetLogger("scheduler", logging.Info) diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go new file mode 100644 index 000000000..2f53804b8 --- /dev/null +++ b/pkg/scheduler/scheduler.go @@ -0,0 +1,119 @@ +// +// DISCLAIMER +// +// Copyright 2024 ArangoDB GmbH, Cologne, Germany +// +// 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. +// +// Copyright holder is ArangoDB GmbH, Cologne, Germany +// + +package scheduler + +import ( + "context" + + core "k8s.io/api/core/v1" + + schedulerApi "github.com/arangodb/kube-arangodb/pkg/apis/scheduler/v1alpha1" + "github.com/arangodb/kube-arangodb/pkg/debug_package/generators/kubernetes" + "github.com/arangodb/kube-arangodb/pkg/util/errors" + "github.com/arangodb/kube-arangodb/pkg/util/kclient" +) + +func NewScheduler(client kclient.Client, namespace string) Scheduler { + return scheduler{ + client: client, + namespace: namespace, + } +} + +type Scheduler interface { + Render(ctx context.Context, in Request, templates ...*schedulerApi.ProfileTemplate) (*core.PodTemplateSpec, []string, error) +} + +type scheduler struct { + client kclient.Client + namespace string +} + +func (s scheduler) Render(ctx context.Context, in Request, templates ...*schedulerApi.ProfileTemplate) (*core.PodTemplateSpec, []string, error) { + profileMap, err := kubernetes.MapObjects[*schedulerApi.ArangoProfileList, *schedulerApi.ArangoProfile](ctx, s.client.Arango().SchedulerV1alpha1().ArangoProfiles(s.namespace), func(result *schedulerApi.ArangoProfileList) []*schedulerApi.ArangoProfile { + q := make([]*schedulerApi.ArangoProfile, len(result.Items)) + + for id, e := range result.Items { + q[id] = e.DeepCopy() + } + + return q + }) + + if err != nil { + return nil, nil, err + } + + profiles := profileMap.AsList().Filter(func(a *schedulerApi.ArangoProfile) bool { + return a != nil && a.Spec.Template != nil + }).Filter(func(a *schedulerApi.ArangoProfile) bool { + if a.Spec.Selectors == nil { + return false + } + + if !a.Spec.Selectors.Select(in.Labels) { + return false + } + + return true + }) + + for _, name := range in.Profiles { + p, ok := profileMap.ByName(name) + if !ok { + return nil, nil, errors.Errorf("Profile with name `%s` is missing", name) + } + + profiles = append(profiles, p) + } + + profiles = profiles.Unique(func(existing kubernetes.List[*schedulerApi.ArangoProfile], o *schedulerApi.ArangoProfile) bool { + return existing.Contains(func(a *schedulerApi.ArangoProfile) bool { + return a.GetName() == o.GetName() + }) + }) + + profiles = profiles.Sort(func(a, b *schedulerApi.ArangoProfile) bool { + return a.Spec.Template.GetPriority() > b.Spec.Template.GetPriority() + }) + + if err := errors.Errors(kubernetes.Extract(profiles, func(in *schedulerApi.ArangoProfile) error { + return in.Spec.Validate() + })...); err != nil { + return nil, nil, err + } + + extracted := schedulerApi.ProfileTemplates(kubernetes.Extract(profiles, func(in *schedulerApi.ArangoProfile) *schedulerApi.ProfileTemplate { + return in.Spec.Template + }).Append(templates...).Append(in.AsTemplate())) + + names := kubernetes.Extract(profiles, func(in *schedulerApi.ArangoProfile) string { + return in.GetName() + }) + + var pod core.PodTemplateSpec + + if err := extracted.RenderOnTemplate(&pod); err != nil { + return nil, names, err + } + + return &pod, names, nil +} diff --git a/pkg/scheduler/scheduler_test.go b/pkg/scheduler/scheduler_test.go new file mode 100644 index 000000000..e41ba0254 --- /dev/null +++ b/pkg/scheduler/scheduler_test.go @@ -0,0 +1,289 @@ +// +// DISCLAIMER +// +// Copyright 2024 ArangoDB GmbH, Cologne, Germany +// +// 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. +// +// Copyright holder is ArangoDB GmbH, Cologne, Germany +// + +package scheduler + +import ( + "context" + "strings" + "testing" + + "github.com/stretchr/testify/require" + core "k8s.io/api/core/v1" + meta "k8s.io/apimachinery/pkg/apis/meta/v1" + "sigs.k8s.io/yaml" + + schedulerApi "github.com/arangodb/kube-arangodb/pkg/apis/scheduler/v1alpha1" + schedulerContainerApi "github.com/arangodb/kube-arangodb/pkg/apis/scheduler/v1alpha1/container" + schedulerContainerResourcesApi "github.com/arangodb/kube-arangodb/pkg/apis/scheduler/v1alpha1/container/resources" + "github.com/arangodb/kube-arangodb/pkg/util" + "github.com/arangodb/kube-arangodb/pkg/util/kclient" + "github.com/arangodb/kube-arangodb/pkg/util/tests" +) + +func newScheduler(t *testing.T, objects ...*schedulerApi.ArangoProfile) Scheduler { + client := kclient.NewFakeClientBuilder().Client() + + objs := make([]interface{}, len(objects)) + for id := range objs { + objs[id] = &objects[id] + } + + tests.CreateObjects(t, client.Kubernetes(), client.Arango(), objs...) + + return NewScheduler(client, tests.FakeNamespace) +} + +type validatorExec func(in validator) + +type validator func(t *testing.T, err error, template *core.PodTemplateSpec, accepted []string) + +func render(t *testing.T, s Scheduler, in Request, templates ...*schedulerApi.ProfileTemplate) validatorExec { + pod, accepted, err := s.Render(context.Background(), in, templates...) + t.Logf("Accepted templates: %s", strings.Join(accepted, ", ")) + if err != nil { + return runValidate(t, err, pod, accepted) + } + require.NoError(t, err) + + data, err := yaml.Marshal(pod) + require.NoError(t, err) + + t.Logf("Rendered Template:\n%s", string(data)) + + return runValidate(t, nil, pod, accepted) +} + +func runValidate(t *testing.T, err error, template *core.PodTemplateSpec, accepted []string) validatorExec { + return func(in validator) { + t.Run("Validate", func(t *testing.T) { + in(t, err, template, accepted) + }) + } +} + +func Test_NoProfiles(t *testing.T) { + render(t, newScheduler(t, tests.NewMetaObjectInDefaultNamespace[*schedulerApi.ArangoProfile](t, "test")), Request{})(func(t *testing.T, err error, template *core.PodTemplateSpec, accepted []string) { + require.NoError(t, err) + + require.Len(t, accepted, 0) + + tests.GetContainerByNameT(t, template.Spec.Containers, DefaultContainerName) + }) +} + +func Test_MissingSelectedProfile(t *testing.T) { + render(t, newScheduler(t, tests.NewMetaObjectInDefaultNamespace[*schedulerApi.ArangoProfile](t, "test")), Request{ + Profiles: []string{"missing"}, + })(func(t *testing.T, err error, template *core.PodTemplateSpec, accepted []string) { + require.EqualError(t, err, "Profile with name `missing` is missing") + }) +} + +func Test_SelectorWithoutSelector(t *testing.T) { + render(t, newScheduler(t, tests.NewMetaObjectInDefaultNamespace[*schedulerApi.ArangoProfile](t, "test", func(t *testing.T, obj *schedulerApi.ArangoProfile) { + obj.Spec.Template = &schedulerApi.ProfileTemplate{ + Container: &schedulerApi.ProfileContainerTemplate{ + Containers: schedulerContainerApi.Containers{ + DefaultContainerName: { + Image: &schedulerContainerResourcesApi.Image{ + Image: util.NewType("image:1"), + }, + }, + }, + }, + } + })), Request{})(func(t *testing.T, err error, template *core.PodTemplateSpec, accepted []string) { + require.NoError(t, err) + + require.Len(t, accepted, 0) + + c := tests.GetContainerByNameT(t, template.Spec.Containers, DefaultContainerName) + require.Equal(t, "", c.Image) + }) +} + +func Test_SelectorWithSelectorAll(t *testing.T) { + render(t, newScheduler(t, tests.NewMetaObjectInDefaultNamespace[*schedulerApi.ArangoProfile](t, "test", func(t *testing.T, obj *schedulerApi.ArangoProfile) { + obj.Spec.Selectors = &schedulerApi.ProfileSelectors{ + Label: &meta.LabelSelector{}, + } + obj.Spec.Template = &schedulerApi.ProfileTemplate{ + Container: &schedulerApi.ProfileContainerTemplate{ + Containers: schedulerContainerApi.Containers{ + DefaultContainerName: { + Image: &schedulerContainerResourcesApi.Image{ + Image: util.NewType("image:1"), + }, + }, + }, + }, + } + })), Request{})(func(t *testing.T, err error, template *core.PodTemplateSpec, accepted []string) { + require.NoError(t, err) + + require.Len(t, accepted, 1) + require.Equal(t, []string{ + "test", + }, accepted) + + c := tests.GetContainerByNameT(t, template.Spec.Containers, DefaultContainerName) + require.Equal(t, "image:1", c.Image) + }) +} + +func Test_SelectorWithSpecificSelector_MissingLabel(t *testing.T) { + render(t, newScheduler(t, tests.NewMetaObjectInDefaultNamespace[*schedulerApi.ArangoProfile](t, "test", func(t *testing.T, obj *schedulerApi.ArangoProfile) { + obj.Spec.Selectors = &schedulerApi.ProfileSelectors{ + Label: &meta.LabelSelector{ + MatchExpressions: []meta.LabelSelectorRequirement{ + { + Key: "ml.arangodb.com/type", + Operator: meta.LabelSelectorOpExists, + }, + }, + }, + } + obj.Spec.Template = &schedulerApi.ProfileTemplate{ + Container: &schedulerApi.ProfileContainerTemplate{ + Containers: schedulerContainerApi.Containers{ + DefaultContainerName: { + Image: &schedulerContainerResourcesApi.Image{ + Image: util.NewType("image:1"), + }, + }, + }, + }, + } + })), Request{})(func(t *testing.T, err error, template *core.PodTemplateSpec, accepted []string) { + require.NoError(t, err) + + require.Len(t, accepted, 0) + + c := tests.GetContainerByNameT(t, template.Spec.Containers, DefaultContainerName) + require.Equal(t, "", c.Image) + }) +} + +func Test_SelectorWithSpecificSelector_PresentLabel(t *testing.T) { + render(t, newScheduler(t, tests.NewMetaObjectInDefaultNamespace[*schedulerApi.ArangoProfile](t, "test", func(t *testing.T, obj *schedulerApi.ArangoProfile) { + obj.Spec.Selectors = &schedulerApi.ProfileSelectors{ + Label: &meta.LabelSelector{ + MatchExpressions: []meta.LabelSelectorRequirement{ + { + Key: "ml.arangodb.com/type", + Operator: meta.LabelSelectorOpExists, + }, + }, + }, + } + obj.Spec.Template = &schedulerApi.ProfileTemplate{ + Container: &schedulerApi.ProfileContainerTemplate{ + Containers: schedulerContainerApi.Containers{ + DefaultContainerName: { + Image: &schedulerContainerResourcesApi.Image{ + Image: util.NewType("image:1"), + }, + }, + }, + }, + } + })), Request{ + Labels: map[string]string{ + "ml.arangodb.com/type": "training", + }, + }, nil)(func(t *testing.T, err error, template *core.PodTemplateSpec, accepted []string) { + require.NoError(t, err) + + require.Len(t, accepted, 1) + require.Equal(t, []string{ + "test", + }, accepted) + + c := tests.GetContainerByNameT(t, template.Spec.Containers, DefaultContainerName) + require.Equal(t, "image:1", c.Image) + }) +} + +func Test_SelectorWithSpecificSelector_PresentLabel_ByPriority(t *testing.T) { + render(t, newScheduler(t, + tests.NewMetaObjectInDefaultNamespace[*schedulerApi.ArangoProfile](t, "test", func(t *testing.T, obj *schedulerApi.ArangoProfile) { + obj.Spec.Selectors = &schedulerApi.ProfileSelectors{ + Label: &meta.LabelSelector{ + MatchExpressions: []meta.LabelSelectorRequirement{ + { + Key: "ml.arangodb.com/type", + Operator: meta.LabelSelectorOpExists, + }, + }, + }, + } + obj.Spec.Template = &schedulerApi.ProfileTemplate{ + Priority: util.NewType(1), + Container: &schedulerApi.ProfileContainerTemplate{ + Containers: schedulerContainerApi.Containers{ + DefaultContainerName: { + Image: &schedulerContainerResourcesApi.Image{ + Image: util.NewType("image:1"), + }, + }, + }, + }, + } + }), tests.NewMetaObjectInDefaultNamespace[*schedulerApi.ArangoProfile](t, "test2", func(t *testing.T, obj *schedulerApi.ArangoProfile) { + obj.Spec.Selectors = &schedulerApi.ProfileSelectors{ + Label: &meta.LabelSelector{ + MatchExpressions: []meta.LabelSelectorRequirement{ + { + Key: "ml.arangodb.com/type", + Operator: meta.LabelSelectorOpExists, + }, + }, + }, + } + obj.Spec.Template = &schedulerApi.ProfileTemplate{ + Priority: util.NewType(2), + Container: &schedulerApi.ProfileContainerTemplate{ + Containers: schedulerContainerApi.Containers{ + DefaultContainerName: { + Image: &schedulerContainerResourcesApi.Image{ + Image: util.NewType("image:2"), + }, + }, + }, + }, + } + })), Request{ + Labels: map[string]string{ + "ml.arangodb.com/type": "training", + }, + })(func(t *testing.T, err error, template *core.PodTemplateSpec, accepted []string) { + require.NoError(t, err) + + require.Len(t, accepted, 2) + require.Equal(t, []string{ + "test2", + "test", + }, accepted) + + c := tests.GetContainerByNameT(t, template.Spec.Containers, DefaultContainerName) + require.Equal(t, "image:2", c.Image) + }) +}