Pytorch v1alpha2 api implementation (#54)

* Pytorcxh v1alpha2 api implementation

* Adding unit tests

Gopkg.toml

@@ -0,0 +1,76 @@
+ignored = ["github.com/kubeflow/pytorch-operator"]
+
+required = [
+  "k8s.io/code-generator/cmd/client-gen",
+  "k8s.io/code-generator/cmd/informer-gen",
+  "k8s.io/code-generator/cmd/lister-gen",
+  "k8s.io/code-generator/cmd/deepcopy-gen",
+  "k8s.io/code-generator/cmd/defaulter-gen",
+  "k8s.io/code-generator/cmd/openapi-gen",
+
+  # needed by generated clientsets, but not an explicit dep in client-gen itself
+  "k8s.io/apimachinery/pkg/apimachinery/registered",
+]
+
+[[constraint]]
+  name = "github.com/ghodss/yaml"
+  version = "1.0.0"
+
+[[constraint]]
+  name = "github.com/gogo/protobuf"
+  version = "1.0.0"
+
+[[constraint]]
+  branch = "master"
+  name = "github.com/golang/glog"
+
+[[constraint]]
+  name = "github.com/golang/protobuf"
+  version = "1.1.0"
+
+[[constraint]]
+  branch = "master"
+  name = "github.com/onrik/logrus"
+
+[[constraint]]
+  name = "github.com/stretchr/testify"
+  version = "1.2.2"
+
+[[constraint]]
+  name = "github.com/sirupsen/logrus"
+  version = "~1.0.4"
+
+[[constraint]]
+  name = "k8s.io/client-go"
+  version = "~6.0.0"
+
+[[constraint]]
+  name = "k8s.io/kubernetes"
+  version = "~v1.9.0"
+
+[[constraint]]
+  name = "k8s.io/api"
+  branch = "release-1.9"
+
+[[constraint]]
+  name = "k8s.io/apimachinery"
+  branch = "release-1.9"
+
+[[constraint]]
+  name = "k8s.io/code-generator"
+  # We can not use master since the generated package name is changed from tensorflow to kubeflow.
+  revision = "961bc077103364eb5bda2c40e2b560d068c9a8c6"
+
+[[override]]
+  # To solve undefined: reference.ParseNormalizedNamed in Kubernetes.
+  name = "github.com/docker/distribution"
+  revision = "edc3ab29cdff8694dd6feb85cfeb4b5f1b38ed9c"
+
+[prune]
+  go-tests = true
+  unused-packages = true
+  non-go = true
+  [[prune.project]]
+    name = "k8s.io/code-generator"
+    unused-packages = false
+    non-go = false

cmd/pytorch-operator.v2/app/options/options.go

@@ -0,0 +1,51 @@
+// Copyright 2018 The Kubeflow 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 (
+	"flag"
+)
+
+// ServerOption is the main context object for the controller manager.
+type ServerOption struct {
+	Kubeconfig           string
+	MasterURL            string
+	Threadiness          int
+	PrintVersion         bool
+	JSONLogFormat        bool
+	EnableGangScheduling bool
+}
+
+// NewServerOption creates a new CMServer with a default config.
+func NewServerOption() *ServerOption {
+	s := ServerOption{}
+	return &s
+}
+
+// AddFlags adds flags for a specific CMServer to the specified FlagSet.
+func (s *ServerOption) AddFlags(fs *flag.FlagSet) {
+	fs.StringVar(&s.MasterURL, "master", "",
+		`The url of the Kubernetes API server,
+		 will overrides any value in kubeconfig, only required if out-of-cluster.`)
+
+	fs.IntVar(&s.Threadiness, "threadiness", 1,
+		`How many threads to process the main logic`)
+
+	fs.BoolVar(&s.PrintVersion, "version", false, "Show version and quit")
+
+	fs.BoolVar(&s.JSONLogFormat, "json-log-format", true,
+		"Set true to use json style log format. Set false to use plaintext style log format")
+	fs.BoolVar(&s.EnableGangScheduling, "enable-gang-scheduling", false, "Set true to enable gang scheduling by kube-arbitrator.")
+}

cmd/pytorch-operator.v2/app/server.go

@@ -0,0 +1,170 @@
+// Copyright 2018 The Kubeflow 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 app
+
+import (
+	"fmt"
+	"os"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	kubeinformers "k8s.io/client-go/informers"
+	kubeclientset "k8s.io/client-go/kubernetes"
+	restclientset "k8s.io/client-go/rest"
+	"k8s.io/client-go/tools/clientcmd"
+	election "k8s.io/client-go/tools/leaderelection"
+	"k8s.io/client-go/tools/leaderelection/resourcelock"
+	"k8s.io/client-go/tools/record"
+
+	"github.com/kubeflow/pytorch-operator/cmd/pytorch-operator.v2/app/options"
+	"github.com/kubeflow/pytorch-operator/pkg/apis/pytorch/v1alpha2"
+	jobclientset "github.com/kubeflow/pytorch-operator/pkg/client/clientset/versioned"
+	"github.com/kubeflow/pytorch-operator/pkg/client/clientset/versioned/scheme"
+	jobinformers "github.com/kubeflow/pytorch-operator/pkg/client/informers/externalversions"
+	controller "github.com/kubeflow/pytorch-operator/pkg/controller.v2/pytorch"
+	"github.com/kubeflow/tf-operator/pkg/util/signals"
+	"github.com/kubeflow/tf-operator/pkg/version"
+)
+
+const (
+	apiVersion = "v1alpha2"
+)
+
+var (
+	// leader election config
+	leaseDuration = 15 * time.Second
+	renewDuration = 5 * time.Second
+	retryPeriod   = 3 * time.Second
+	resyncPeriod  = 30 * time.Second
+)
+
+const RecommendedKubeConfigPathEnv = "KUBECONFIG"
+
+func Run(opt *options.ServerOption) error {
+	// Check if the -version flag was passed and, if so, print the version and exit.
+	if opt.PrintVersion {
+		version.PrintVersionAndExit(apiVersion)
+	}
+
+	namespace := os.Getenv(v1alpha2.EnvKubeflowNamespace)
+	if len(namespace) == 0 {
+		log.Infof("EnvKubeflowNamespace not set, use default namespace")
+		namespace = metav1.NamespaceDefault
+	}
+
+	// To help debugging, immediately log version.
+	log.Infof("%+v", version.Info(apiVersion))
+
+	// Set up signals so we handle the first shutdown signal gracefully.
+	stopCh := signals.SetupSignalHandler()
+
+	// Note: ENV KUBECONFIG will overwrite user defined Kubeconfig option.
+	if len(os.Getenv(RecommendedKubeConfigPathEnv)) > 0 {
+		// use the current context in kubeconfig
+		// This is very useful for running locally.
+		opt.Kubeconfig = os.Getenv(RecommendedKubeConfigPathEnv)
+	}
+
+	// Get kubernetes config.
+	kcfg, err := clientcmd.BuildConfigFromFlags(opt.MasterURL, opt.Kubeconfig)
+	if err != nil {
+		log.Fatalf("Error building kubeconfig: %s", err.Error())
+	}
+
+	// Create clients.
+	kubeClientSet, leaderElectionClientSet, pytorchJobClientSet, err := createClientSets(kcfg)
+	if err != nil {
+		return err
+	}
+
+	// Create informer factory.
+	kubeInformerFactory := kubeinformers.NewSharedInformerFactory(kubeClientSet, resyncPeriod)
+	pytorchJobInformerFactory := jobinformers.NewSharedInformerFactory(pytorchJobClientSet, resyncPeriod)
+
+	unstructuredInformer := controller.NewUnstructuredPyTorchJobInformer(kcfg)
+
+	// Create pytorch controller.
+	tc := controller.NewPyTorchController(unstructuredInformer, kubeClientSet, pytorchJobClientSet, kubeInformerFactory, pytorchJobInformerFactory, *opt)
+
+	// Start informer goroutines.
+	go kubeInformerFactory.Start(stopCh)
+
+	go unstructuredInformer.Informer().Run(stopCh)
+
+	// Set leader election start function.
+	run := func(<-chan struct{}) {
+		if err := tc.Run(opt.Threadiness, stopCh); err != nil {
+			log.Errorf("Failed to run the controller: %v", err)
+		}
+	}
+
+	id, err := os.Hostname()
+	if err != nil {
+		return fmt.Errorf("Failed to get hostname: %v", err)
+	}
+
+	// Prepare event clients.
+	eventBroadcaster := record.NewBroadcaster()
+	recorder := eventBroadcaster.NewRecorder(scheme.Scheme, v1.EventSource{Component: "pytorch-operator"})
+
+	rl := &resourcelock.EndpointsLock{
+		EndpointsMeta: metav1.ObjectMeta{
+			Namespace: namespace,
+			Name:      "pytorch-operator",
+		},
+		Client: leaderElectionClientSet.CoreV1(),
+		LockConfig: resourcelock.ResourceLockConfig{
+			Identity:      id,
+			EventRecorder: recorder,
+		},
+	}
+
+	// Start leader election.
+	election.RunOrDie(election.LeaderElectionConfig{
+		Lock:          rl,
+		LeaseDuration: leaseDuration,
+		RenewDeadline: renewDuration,
+		RetryPeriod:   retryPeriod,
+		Callbacks: election.LeaderCallbacks{
+			OnStartedLeading: run,
+			OnStoppedLeading: func() {
+				log.Fatalf("leader election lost")
+			},
+		},
+	})
+
+	return nil
+}
+
+func createClientSets(config *restclientset.Config) (kubeclientset.Interface, kubeclientset.Interface, jobclientset.Interface, error) {
+	kubeClientSet, err := kubeclientset.NewForConfig(restclientset.AddUserAgent(config, "pytorch-operator"))
+	if err != nil {
+		return nil, nil, nil, err
+	}
+
+	leaderElectionClientSet, err := kubeclientset.NewForConfig(restclientset.AddUserAgent(config, "leader-election"))
+	if err != nil {
+		return nil, nil, nil, err
+	}
+
+	jobClientSet, err := jobclientset.NewForConfig(config)
+	if err != nil {
+		return nil, nil, nil, err
+	}
+
+	return kubeClientSet, leaderElectionClientSet, jobClientSet, nil
+}

cmd/pytorch-operator.v2/main.go

@@ -0,0 +1,49 @@
+// Copyright 2018 The Kubeflow 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 main
+
+import (
+	"flag"
+
+	"github.com/onrik/logrus/filename"
+	log "github.com/sirupsen/logrus"
+
+	"github.com/kubeflow/pytorch-operator/cmd/pytorch-operator.v2/app"
+	"github.com/kubeflow/pytorch-operator/cmd/pytorch-operator.v2/app/options"
+)
+
+func init() {
+	// Add filename as one of the fields of the structured log message.
+	filenameHook := filename.NewHook()
+	filenameHook.Field = "filename"
+	log.AddHook(filenameHook)
+}
+
+func main() {
+	s := options.NewServerOption()
+	s.AddFlags(flag.CommandLine)
+
+	flag.Parse()
+
+	if s.JSONLogFormat {
+		// Output logs in a json format so that it can be parsed by services like Stackdriver.
+		log.SetFormatter(&log.JSONFormatter{})
+	}
+
+	if err := app.Run(s); err != nil {
+		log.Fatalf("%v\n", err)
+	}
+
+}