sedna/pkg/globalmanager/runtime/worker.go

package runtime

import (
	"context"
	"fmt"
	sednav1 "github.com/kubeedge/sedna/pkg/apis/sedna/v1alpha1"
	autoscalingv2 "k8s.io/api/autoscaling/v2"
	"k8s.io/apimachinery/pkg/api/errors"
	"path/filepath"
	"strconv"
	"strings"

	appsv1 "k8s.io/api/apps/v1"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/client-go/kubernetes"
	"k8s.io/klog/v2"

	"github.com/kubeedge/sedna/pkg/globalmanager/utils"
)

type WorkerMount struct {
	Name string
	// the url to be mounted
	URL *MountURL

	// for some cases, there are more than one url to be mounted
	URLs []MountURL

	// envName indicates the environment key of the mounts injected to the worker
	EnvName string
}

// WorkerParam describes the system-defined parameters of worker
type WorkerParam struct {
	Mounts []WorkerMount

	Env        map[string]string
	WorkerType string

	// if true, force to use hostNetwork
	HostNetwork bool

	ModelHotUpdate ModelHotUpdate

	RestartPolicy v1.RestartPolicy

	DNSPolicy v1.DNSPolicy
}

type ModelHotUpdate struct {
	Enable            bool
	PollPeriodSeconds int64
}

// generateLabels generates labels for an object
func generateLabels(object CommonInterface, workerType string) map[string]string {
	kind := object.GroupVersionKind().Kind
	group := object.GroupVersionKind().Group

	keyPrefix := strings.ToLower(kind + "." + group + "/")

	labels := make(map[string]string)
	labels[keyPrefix+"name"] = object.GetName()
	labels[keyPrefix+"uid"] = string(object.GetUID())
	if workerType != "" {
		labels[keyPrefix+"worker-type"] = strings.ToLower(workerType)
	}
	return labels
}

// GenerateSelector generates the selector of an object for worker
func GenerateSelector(object CommonInterface) (labels.Selector, error) {
	ls := &metav1.LabelSelector{
		// select any type workers
		MatchLabels: generateLabels(object, ""),
	}
	return metav1.LabelSelectorAsSelector(ls)
}

// GenerateWorkerSelector generates the selector of an object for specific worker type
func GenerateWorkerSelector(object CommonInterface, workerType string) (labels.Selector, error) {
	ls := &metav1.LabelSelector{
		// select any type workers
		MatchLabels: generateLabels(object, workerType),
	}
	return metav1.LabelSelectorAsSelector(ls)
}

// CreateKubernetesService creates a k8s service for an object given ip and port
func CreateKubernetesService(kubeClient kubernetes.Interface, object CommonInterface, workerType string, inputPort int32, inputIP string) (int32, error) {
	ctx := context.Background()
	name := object.GetName()
	namespace := object.GetNamespace()
	kind := object.GroupVersionKind().Kind
	targePort := intstr.IntOrString{
		IntVal: inputPort,
	}
	serviceSpec := &v1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: object.GetNamespace(),
			Name:      strings.ToLower(name + "-" + workerType),
			OwnerReferences: []metav1.OwnerReference{
				*metav1.NewControllerRef(object, object.GroupVersionKind()),
			},
			Labels: generateLabels(object, workerType),
		},
		Spec: v1.ServiceSpec{
			Selector: generateLabels(object, workerType),
			ExternalIPs: []string{
				inputIP,
			},
			Type: v1.ServiceTypeNodePort,
			Ports: []v1.ServicePort{
				{
					Port:       inputPort,
					TargetPort: targePort,
				},
			},
		},
	}
	service, err := kubeClient.CoreV1().Services(namespace).Create(ctx, serviceSpec, metav1.CreateOptions{})
	if err != nil {
		klog.Warningf("failed to create service for %v %v/%v, err:%s", kind, namespace, name, err)
		return 0, err
	}

	klog.V(2).Infof("Service %s is created successfully for %v %v/%v", service.Name, kind, namespace, name)
	return service.Spec.Ports[0].NodePort, nil
}

// injectWorkerParam modifies pod in-place
func injectWorkerParam(pod *v1.Pod, workerParam *WorkerParam, object CommonInterface) {
	InjectStorageInitializer(pod, workerParam)

	if workerParam.WorkerType == InferencePodType && workerParam.ModelHotUpdate.Enable {
		injectModelHotUpdateMount(pod, object)
		setModelHotUpdateEnv(workerParam)
	}

	envs := createEnvVars(workerParam.Env)
	for idx := range pod.Spec.Containers {
		pod.Spec.Containers[idx].Env = append(
			pod.Spec.Containers[idx].Env, envs...,
		)
	}

	// inject our labels
	if pod.Labels == nil {
		pod.Labels = make(map[string]string)
	}

	for k, v := range generateLabels(object, workerParam.WorkerType) {
		pod.Labels[k] = v
	}

	pod.GenerateName = object.GetName() + "-" + strings.ToLower(workerParam.WorkerType) + "-"

	pod.Namespace = object.GetNamespace()

	if workerParam.HostNetwork {
		// FIXME
		// force to set hostnetwork
		pod.Spec.HostNetwork = true
	}

	if pod.Spec.RestartPolicy == "" {
		pod.Spec.RestartPolicy = workerParam.RestartPolicy
	}

	if workerParam.DNSPolicy != "" {
		pod.Spec.DNSPolicy = workerParam.DNSPolicy
	}
}

// CreatePodWithTemplate creates and returns a pod object given a crd object, pod template, and workerParam
func CreatePodWithTemplate(client kubernetes.Interface, object CommonInterface, spec *v1.PodTemplateSpec, workerParam *WorkerParam) (*v1.Pod, error) {
	objectKind := object.GroupVersionKind()
	pod, _ := utils.GetPodFromTemplate(spec, object, metav1.NewControllerRef(object, objectKind))
	injectWorkerParam(pod, workerParam, object)

	createdPod, err := client.CoreV1().Pods(object.GetNamespace()).Create(context.TODO(), pod, metav1.CreateOptions{})
	objectName := object.GetNamespace() + "/" + object.GetName()
	if err != nil {
		klog.Warningf("failed to create pod(type=%s) for %s %s, err:%s", workerParam.WorkerType, objectKind, objectName, err)
		return nil, err
	}
	klog.V(2).Infof("pod %s is created successfully for %s %s", createdPod.Name, objectKind, objectName)
	return createdPod, nil
}

// CreateEdgeMeshService creates a kubeedge edgemesh service for an object, and returns an edgemesh service URL.
// Since edgemesh can realize Cross-Edge-Cloud communication, the service can be created both on the cloud or edge side.
func CreateEdgeMeshService(kubeClient kubernetes.Interface, object CommonInterface, workerType string, servicePort int32) (string, error) {
	ctx := context.Background()
	name := object.GetName()
	namespace := object.GetNamespace()
	kind := object.GroupVersionKind().Kind
	targetPort := intstr.IntOrString{
		IntVal: servicePort,
	}
	serviceSpec := &v1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: namespace,
			Name:      strings.ToLower(name + "-" + workerType),
			OwnerReferences: []metav1.OwnerReference{
				*metav1.NewControllerRef(object, object.GroupVersionKind()),
			},
			Labels: generateLabels(object, workerType),
		},
		Spec: v1.ServiceSpec{
			Selector: generateLabels(object, workerType),
			Ports: []v1.ServicePort{
				{
					// TODO: be clean, Port.Name is currently required by edgemesh(v1.8.0).
					// and should be <protocol>-<suffix>
					Name:       "tcp-0",
					Protocol:   "TCP",
					Port:       servicePort,
					TargetPort: targetPort,
				},
			},
		},
	}
	service, err := kubeClient.CoreV1().Services(namespace).Create(ctx, serviceSpec, metav1.CreateOptions{})
	if err != nil {
		klog.Warningf("failed to create service for %v %v/%v, err:%s", kind, namespace, name, err)
		return "", err
	}

	klog.V(2).Infof("Service %s is created successfully for %v %v/%v", service.Name, kind, namespace, name)
	return fmt.Sprintf("%s.%s", service.Name, service.Namespace), nil
}

// CreateDeploymentWithTemplate creates and returns a deployment object given a crd object, deployment template
func CreateDeploymentWithTemplate(client kubernetes.Interface, object CommonInterface, spec *appsv1.DeploymentSpec, workerParam *WorkerParam) (*appsv1.Deployment, error) {
	objectKind := object.GroupVersionKind()
	objectName := object.GetNamespace() + "/" + object.GetName()
	deployment := newDeployment(object, spec, workerParam)

	injectDeploymentParam(deployment, workerParam, object)

	createdDeployment, err := client.AppsV1().Deployments(object.GetNamespace()).Create(context.TODO(), deployment, metav1.CreateOptions{})
	if err != nil {
		klog.Warningf("failed to create deployment for %s %s, err:%s", objectKind, objectName, err)
		return nil, err
	}
	klog.V(2).Infof("deployment %s is created successfully for %s %s", createdDeployment.Name, objectKind, objectName)
	return createdDeployment, nil
}

// UpdateDeploymentWithTemplate updates an existing deployment object given a crd object, deployment template, and worker parameters
func UpdateDeploymentWithTemplate(client kubernetes.Interface, object CommonInterface, newDeployment *appsv1.Deployment, workerParam *WorkerParam) (*appsv1.Deployment, error) {
	objectKind := object.GroupVersionKind()
	objectName := object.GetNamespace() + "/" + object.GetName()

	// Inject worker parameters.
	injectDeploymentParam(newDeployment, workerParam, object)

	// Call the Kubernetes API to perform the update.
	updatedDeployment, err := client.AppsV1().Deployments(newDeployment.Namespace).Update(context.TODO(), newDeployment, metav1.UpdateOptions{})
	if err != nil {
		klog.Warningf("failed to update deployment for %s %s, err: %s", objectKind, objectName, err)
		return nil, fmt.Errorf("failed to update deployment: %w", err)
	}

	klog.V(2).Infof("deployment %s is updated successfully for %s %s", updatedDeployment.Name, objectKind, objectName)
	return updatedDeployment, nil
}

func CreateHPA(client kubernetes.Interface, object CommonInterface, kind, scaleTargetRefName, workerType string, hpa *sednav1.HPA) error {
	hpaName := "hpa-" + scaleTargetRefName
	newHPA := &autoscalingv2.HorizontalPodAutoscaler{
		ObjectMeta: metav1.ObjectMeta{
			Name:      hpaName,
			Namespace: object.GetNamespace(),
			OwnerReferences: []metav1.OwnerReference{
				*metav1.NewControllerRef(object, object.GroupVersionKind()),
			},
			Labels: generateLabels(object, workerType),
		},
		Spec: autoscalingv2.HorizontalPodAutoscalerSpec{
			MaxReplicas: hpa.MaxReplicas,
			Metrics:     hpa.Metrics,
			MinReplicas: hpa.MinReplicas,
			ScaleTargetRef: autoscalingv2.CrossVersionObjectReference{
				APIVersion: "apps/v1",
				Kind:       kind,
				Name:       scaleTargetRefName,
			},
			Behavior: hpa.Behavior,
		},
	}
	_, err := client.AutoscalingV2().HorizontalPodAutoscalers(object.GetNamespace()).Create(context.TODO(), newHPA, metav1.CreateOptions{})
	if err != nil {
		return fmt.Errorf("failed to create hpa for %s %s, err: %s", kind, hpaName, err)
	}
	return nil
}

func UpdateHPA(client kubernetes.Interface, object CommonInterface, kind, scaleTargetRefName, workerType string, hpa *sednav1.HPA) error {
	// get existing HPA
	hpaName := "hpa-" + scaleTargetRefName
	existingHPA, err := client.AutoscalingV2().HorizontalPodAutoscalers(object.GetNamespace()).Get(context.TODO(), hpaName, metav1.GetOptions{})
	if err != nil {
		// create HPA if not found
		if errors.IsNotFound(err) {
			klog.Info("hpa not found, creating new hpa...")
			return CreateHPA(client, object, kind, scaleTargetRefName, workerType, hpa)
		}
		return fmt.Errorf("failed to get hpa for %s %s, err: %s", kind, hpaName, err)
	}

	// update HPA
	existingHPA.ObjectMeta.Labels = generateLabels(object, workerType)
	existingHPA.ObjectMeta.OwnerReferences = []metav1.OwnerReference{
		*metav1.NewControllerRef(object, object.GroupVersionKind()),
	}
	existingHPA.Spec.MaxReplicas = hpa.MaxReplicas
	existingHPA.Spec.MinReplicas = hpa.MinReplicas
	existingHPA.Spec.Metrics = hpa.Metrics
	existingHPA.Spec.ScaleTargetRef = autoscalingv2.CrossVersionObjectReference{
		APIVersion: "apps/v1",
		Kind:       kind,
		Name:       scaleTargetRefName,
	}
	existingHPA.Spec.Behavior = hpa.Behavior

	// update HPA
	_, err = client.AutoscalingV2().HorizontalPodAutoscalers(object.GetNamespace()).Update(context.TODO(), existingHPA, metav1.UpdateOptions{})
	if err != nil {
		return fmt.Errorf("failed to update hpa for %s %s, err: %s", kind, hpaName, err)
	}

	return nil
}

func DeleteHPA(client kubernetes.Interface, namespace, name string) error {
	// check if HPA exists
	_, err := client.AutoscalingV2().HorizontalPodAutoscalers(namespace).Get(context.TODO(), name, metav1.GetOptions{})
	if err != nil {
		// Return nil if HPA not found
		if errors.IsNotFound(err) {
			return nil
		}

		return fmt.Errorf("failed to get hpa %s in namespace %s, err: %s", name, namespace, err)
	}

	// delete HPA
	err = client.AutoscalingV2().HorizontalPodAutoscalers(namespace).Delete(context.TODO(), name, metav1.DeleteOptions{})
	if err != nil {
		return fmt.Errorf("failed to delete hpa %s in namespace %s, err: %s", name, namespace, err)
	}

	return nil
}

func newDeployment(object CommonInterface, spec *appsv1.DeploymentSpec, workerParam *WorkerParam) *appsv1.Deployment {
	nameSpace := object.GetNamespace()
	deploymentName := object.GetName() + "-" + "deployment" + "-" + strings.ToLower(workerParam.WorkerType)
	matchLabel := make(map[string]string)
	return &appsv1.Deployment{
		ObjectMeta: metav1.ObjectMeta{
			Name:      deploymentName,
			Namespace: nameSpace,
			OwnerReferences: []metav1.OwnerReference{
				*metav1.NewControllerRef(object, object.GroupVersionKind()),
			},
		},
		Spec: appsv1.DeploymentSpec{
			Replicas: (*spec).Replicas,
			Template: (*spec).Template,
			Selector: &metav1.LabelSelector{
				MatchLabels: matchLabel,
			},
		},
	}
}

// injectDeploymentParam modifies deployment in-place
func injectDeploymentParam(deployment *appsv1.Deployment, workerParam *WorkerParam, object CommonInterface) {
	var appLabelKey = "app.sedna.io"
	var appLabelValue = object.GetName() + "-" + workerParam.WorkerType + "-" + "svc"

	// Injection of the storage variables must be done before loading
	// the environment variables!
	if workerParam.Mounts != nil {
		InjectStorageInitializerDeployment(deployment, workerParam)
	}

	// inject our labels
	if deployment.Labels == nil {
		deployment.Labels = make(map[string]string)
	}
	if deployment.Spec.Template.Labels == nil {
		deployment.Spec.Template.Labels = make(map[string]string)
	}
	if deployment.Spec.Selector.MatchLabels == nil {
		deployment.Spec.Selector.MatchLabels = make(map[string]string)
	}

	for k, v := range generateLabels(object, workerParam.WorkerType) {
		deployment.Labels[k] = v
		deployment.Spec.Template.Labels[k] = v
		deployment.Spec.Selector.MatchLabels[k] = v
	}

	// Edgemesh part, useful for service mapping (not necessary!)
	deployment.Labels[appLabelKey] = appLabelValue
	deployment.Spec.Template.Labels[appLabelKey] = appLabelValue
	deployment.Spec.Selector.MatchLabels[appLabelKey] = appLabelValue

	// Env variables injection
	envs := createEnvVars(workerParam.Env)
	for idx := range deployment.Spec.Template.Spec.Containers {
		deployment.Spec.Template.Spec.Containers[idx].Env = append(
			deployment.Spec.Template.Spec.Containers[idx].Env, envs...,
		)
	}
}

// createEnvVars creates EnvMap for container
// include EnvName and EnvValue map for stage of creating a pod
func createEnvVars(envMap map[string]string) []v1.EnvVar {
	var envVars []v1.EnvVar
	for envName, envValue := range envMap {
		Env := v1.EnvVar{
			Name:  envName,
			Value: envValue,
		}
		envVars = append(envVars, Env)
	}
	return envVars
}

// injectModelHotUpdateMount injects volume mounts when worker supports hot update of model
func injectModelHotUpdateMount(pod *v1.Pod, object CommonInterface) {
	hostPathType := v1.HostPathDirectoryOrCreate

	var volumes []v1.Volume
	var volumeMounts []v1.VolumeMount

	modelHotUpdateHostDir, _ := filepath.Split(GetModelHotUpdateConfigFile(object, ModelHotUpdateHostPrefix))
	volumeName := ConvertK8SValidName(ModelHotUpdateVolumeName)
	volumes = append(volumes, v1.Volume{
		Name: volumeName,
		VolumeSource: v1.VolumeSource{
			HostPath: &v1.HostPathVolumeSource{
				Path: modelHotUpdateHostDir,
				Type: &hostPathType,
			},
		},
	})

	volumeMounts = append(volumeMounts, v1.VolumeMount{
		MountPath: ModelHotUpdateContainerPrefix,
		Name:      volumeName,
	})

	injectVolume(pod, volumes, volumeMounts)
}

func GetModelHotUpdateConfigFile(object CommonInterface, prefix string) string {
	return strings.ToLower(filepath.Join(prefix, object.GetNamespace(), object.GetObjectKind().GroupVersionKind().Kind,
		object.GetName(), ModelHotUpdateConfigFile))
}

// setModelHotUpdateEnv sets envs of model hot update
func setModelHotUpdateEnv(workerParam *WorkerParam) {
	workerParam.Env["MODEL_HOT_UPDATE"] = "true"
	workerParam.Env["MODEL_POLL_PERIOD_SECONDS"] = strconv.FormatInt(workerParam.ModelHotUpdate.PollPeriodSeconds, 10)
	workerParam.Env["MODEL_HOT_UPDATE_CONFIG"] = filepath.Join(ModelHotUpdateContainerPrefix, ModelHotUpdateConfigFile)
}