run.go

// Copyright © 2017 RooFoods LTD
// 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 steps

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"time"

	"github.com/spf13/cobra"
	v1 "k8s.io/api/core/v1"
	k8errors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apimachinery/pkg/util/yaml"
	"k8s.io/client-go/kubernetes"

	"github.com/deliveroo/paddle/cli/pipeline"
)

type runCmdFlagsStruct struct {
	PipelineName       string
	StepJSON           string
	BucketName         string
	SnsArn             string
	Namespace          string
	RunIdentifier      string
	ImageTag           string
	StepBranch         string
	StepVersion        string
	OverrideInputs     bool
	TailLogs           bool
	Secrets            []string
	Env                []string
	BucketOverrides    []string
	DeletePollInterval time.Duration
	StartTimeout       time.Duration
}

const defaultDeletePollInterval = 2 * time.Second
const deleteTimeout = 120 * time.Second
const defaultStartTimeout = 10 * time.Minute

var runCmdFlags *runCmdFlagsStruct
var clientset kubernetes.Interface

var logFatalf = log.Fatalf

var runCmd = &cobra.Command{
	Use:   "run [pipeline_yaml]",
	Short: "Run a pipeline or a pipeline step",
	Args:  cobra.ExactArgs(1),
	Long: `Run a pipeline (or a single step) on the Kubernetes cluster.

Example:

$ paddle steps run -s [step_json]
`,
	Run: func(cmd *cobra.Command, args []string) {
		runStep(runCmdFlags)
	},
}

func init() {
	runCmdFlags = &runCmdFlagsStruct{}
	runCmd.Flags().StringVarP(&runCmdFlags.PipelineName, "pipeline", "p", "", "Pipeline name")
	runCmd.Flags().StringVarP(&runCmdFlags.RunIdentifier, "identifier", "i", "", "Run identifier")
	runCmd.Flags().StringVarP(&runCmdFlags.StepJSON, "step", "s", "", "JSON of the step to execute")
	runCmd.Flags().StringVarP(&runCmdFlags.Namespace, "namespace", "n", "modeltraining", "Kubernetes namespace")
	runCmd.Flags().StringVarP(&runCmdFlags.BucketName, "bucket", "b", "", "Bucket name")
	runCmd.Flags().StringVarP(&runCmdFlags.SnsArn, "sns", "", "", "SNS topic ARN")
	runCmd.Flags().StringVarP(&runCmdFlags.ImageTag, "tag", "t", "", "Image tag (overrides the one defined in the pipeline)")
	runCmd.Flags().StringVarP(&runCmdFlags.StepBranch, "step-branch", "B", "", "Step branch (overrides the one defined in the pipeline)")
	runCmd.Flags().StringVarP(&runCmdFlags.StepVersion, "step-version", "V", "", "Step version (overrides the one defined in the pipeline)")
	runCmd.Flags().BoolVarP(&runCmdFlags.TailLogs, "logs", "l", true, "Tail logs")
	runCmd.Flags().BoolVarP(&runCmdFlags.OverrideInputs, "override-inputs", "I", false, "Override input version/branch (only makes sense to use with -B or -V)")
	runCmd.Flags().StringSliceVarP(&runCmdFlags.Secrets, "secret", "S", []string{}, "Secret to pull into the environment (in the form ENV_VAR:secret_store:key_name)")
	runCmd.Flags().StringSliceVarP(&runCmdFlags.Env, "env", "e", []string{}, "Environment variables to set (in the form name:value)")
	runCmd.Flags().StringSliceVar(&runCmdFlags.BucketOverrides, "replace-input-buckets", []string{}, "Override input bucket names (in the form original_bucket_name:new_bucket_name)")
	runCmdFlags.DeletePollInterval = defaultDeletePollInterval
	runCmdFlags.StartTimeout = defaultStartTimeout

	config, err := pipeline.GetKubernetesConfig()
	if err != nil {
		panic(err.Error())
	}
	clientset, err = kubernetes.NewForConfig(config)
	if err != nil {
		panic(err.Error())
	}
}

func runStep(flags *runCmdFlagsStruct) {

	definition := pipeline.PipelineDefinition{
		Pipeline:  flags.PipelineName,
		Bucket:    flags.BucketName,
		Namespace: flags.Namespace,
	}

	step := pipeline.PipelineDefinitionStep{}

	err := json.Unmarshal([]byte(flags.StepJSON), &step)
	if err != nil {
		log.Fatalf("error: %v", err)
	}

	if flags.ImageTag != "" {
		step.OverrideTag(flags.ImageTag)
	}
	if flags.StepBranch != "" {
		step.OverrideBranch(flags.StepBranch, flags.OverrideInputs)
	}
	if flags.StepVersion != "" {
		step.OverrideVersion(flags.StepVersion, flags.OverrideInputs)
	}
	err = runPipelineStep(&definition, &step, flags)
	if err != nil {
		logFatalf("[paddle] %s", err.Error())
	}
}

func runPipelineStep(definition *pipeline.PipelineDefinition, step *pipeline.PipelineDefinitionStep, flags *runCmdFlagsStruct) error {
	log.Printf("[paddle] Running step %s", step.Step)
	podDefinition := NewPodDefinition(definition, step, flags.SnsArn)
	podDefinition.RunIdentifier = flags.RunIdentifier
	podDefinition.parseSecrets(flags.Secrets)
	podDefinition.parseEnv(flags.Env)
	podDefinition.setBucketOverrides(flags.BucketOverrides)

	stepPodBuffer := podDefinition.compile()
	pod := &v1.Pod{}
	err := yaml.NewYAMLOrJSONDecoder(stepPodBuffer, 4096).Decode(pod)
	if err != nil {
		return err
	}

	pods := clientset.CoreV1().Pods(definition.Namespace)

	err = deleteAndWait(clientset, podDefinition, flags)
	if err != nil {
		return err
	}

	if podDefinition.needsVolume() {
		err := createVolumeClaim(clientset, podDefinition, flags)
		if err != nil {
			return err
		}
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	watch, err := pipeline.Watch(ctx, clientset, pod)
	if err != nil {
		return err
	}

	pod, err = pods.Create(pod)
	if err != nil {
		return err
	}

	containers := make(map[string]bool)

	go func() {
		time.Sleep(flags.StartTimeout)
		if len(containers) < len(pod.Spec.Containers) {
			cancel()
		}
	}()

	removed := map[string]bool{}

	for {
		select {
		case e := <-watch:
			switch e.Type {
			case pipeline.Added:
				log.Printf("[paddle] Container %s/%s starting", pod.Name, e.Container)
				containers[e.Container] = true
				if flags.TailLogs {
					pipeline.TailLogs(ctx, clientset, e.Pod, e.Container)
				}
			case pipeline.Deleted:
				log.Println("[paddle] Pod deleted")
				return errors.New("Pod was deleted unexpectedly.")
			case pipeline.Removed:
				if !removed[e.Container] {
					log.Printf("[paddle] Container removed: %s", e.Container)
				}
				removed[e.Container] = true
				continue
			case pipeline.Completed:
				log.Printf("[paddle] Pod execution completed")
				if podDefinition.needsVolume() {
					deleteVolumeClaim(clientset, podDefinition, flags)
				}
				return nil
			case pipeline.Failed:
				var msg string
				if e.Container != "" {
					if e.Message != "" {
						msg = fmt.Sprintf("Container %s/%s failed: '%s'", pod.Name, e.Container, e.Message)
					} else {
						msg = fmt.Sprintf("Container %s/%s failed", pod.Name, e.Container)
					}
					_, present := containers[e.Container]
					if !present && flags.TailLogs { // container died before being added
						pipeline.TailLogs(ctx, clientset, e.Pod, e.Container)
						time.Sleep(3 * time.Second) // give it time to tail logs
					}
				} else {
					msg = "Pod failed"
				}
				return errors.New(msg)
			}
		case <-ctx.Done():
			pod, _ := pods.Get(podDefinition.PodName, metav1.GetOptions{})
			reason := "Timed out waiting for pod to start. Cluster might not have sufficient resources."
			if pod != nil {
				for _, container := range pod.Status.ContainerStatuses {
					if container.State.Waiting != nil {
						reason = container.State.Waiting.Message
					}
				}
			}
			pods.Delete(podDefinition.PodName, &metav1.DeleteOptions{})
			return errors.New(reason)
		}
	}

	log.Printf("[paddle] Finishing pod execution")
	return nil
}

func deleteAndWait(c kubernetes.Interface, podDefinition *PodDefinition, flags *runCmdFlagsStruct) error {
	pods := clientset.CoreV1().Pods(podDefinition.Namespace)
	deleting := false
	var gracePeriod int64
	opts := metav1.DeleteOptions{GracePeriodSeconds: &gracePeriod}
	err := wait.PollImmediate(flags.DeletePollInterval, deleteTimeout, func() (bool, error) {
		var err error
		err = pods.Delete(podDefinition.PodName, &opts)
		if err != nil {
			if k8errors.IsNotFound(err) {
				if deleting {
					log.Printf("[paddle] deleted pod %s", podDefinition.PodName)
				}
				return true, nil
			} else {
				return true, err
			}
		}
		if !deleting {
			log.Printf("[paddle] deleting pod %s", podDefinition.PodName)
			deleting = true
		}
		return false, nil
	})
	return err
}

func createVolumeClaim(c kubernetes.Interface, podDefinition *PodDefinition, flags *runCmdFlagsStruct) error {
	err := deleteVolumeClaim(c, podDefinition, flags)
	if err != nil {
		return err
	}

	log.Printf("[paddle] Creating volume claim for %s", podDefinition.PodName)
	claim := &v1.PersistentVolumeClaim{}
	claimBuffer := podDefinition.compileVolumeClaim()
	yaml.NewYAMLOrJSONDecoder(claimBuffer, 4096).Decode(claim)

	claims := clientset.CoreV1().PersistentVolumeClaims(podDefinition.Namespace)

	_, err = claims.Create(claim)
	if err != nil {
		return err
	}
	log.Printf("[paddle] Created volume claim %s", claim.Name)
	return nil
}

func deleteVolumeClaim(c kubernetes.Interface, podDefinition *PodDefinition, flags *runCmdFlagsStruct) error {
	claim := &v1.PersistentVolumeClaim{}
	claimBuffer := podDefinition.compileVolumeClaim()
	yaml.NewYAMLOrJSONDecoder(claimBuffer, 4096).Decode(claim)

	claims := clientset.CoreV1().PersistentVolumeClaims(podDefinition.Namespace)

	deleting := false
	var gracePeriod int64
	opts := metav1.DeleteOptions{GracePeriodSeconds: &gracePeriod}
	err := wait.PollImmediate(flags.DeletePollInterval, deleteTimeout, func() (bool, error) {
		var err error
		err = claims.Delete(claim.Name, &opts)
		if err != nil {
			if k8errors.IsNotFound(err) {
				if deleting {
					log.Printf("[paddle] Deleted volume claim %s", claim.Name)
				}
				return true, nil
			} else {
				return true, err
			}
		}
		if !deleting {
			log.Printf("[paddle] Deleting volume claim %s", claim.Name)
			deleting = true
		}
		return false, nil
	})
	if err != nil {
		return err
	}
	return nil
}