Adding Pytorch v1beta1 image (kubeflow#1978)

StatCan · Nov 22, 2018 · a2f93b2 · a2f93b2
1 parent 45fc8aa
commit a2f93b2
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 141 deletions.
diff --git a/kubeflow/pytorch-job/prototypes/pytorch-job.jsonnet b/kubeflow/pytorch-job/prototypes/pytorch-job.jsonnet
@@ -8,44 +8,12 @@
 // @optionalParam numMasters number 1 The number of masters to use
 // @optionalParam numWorkers number 1 The number of workers to use
 // @optionalParam numGpus number 0 The number of GPUs to attach to workers.
-// @optionalParam jobVersion string v1alpha2 The pytorch operator version to use
+// @optionalParam jobVersion string v1beta1 The pytorch operator version to use
 
 local k = import "k.libsonnet";
 
 local util = {
   pytorchJobReplica(replica_type, number, args, image, num_gpus=0)::
-    local baseContainer = {
-      image: image,
-      name: "pytorch",
-    };
-    local containerArgs = if std.length(args) > 0 then
-      {
-        args: args,
-      }
-    else {};
-    local resources = if num_gpus > 0 then {
-      resources: {
-        limits: {
-          "nvidia.com/gpu": num_gpus,
-        },
-      },
-    } else {};
-    if number > 0 then
-      {
-        replicas: number,
-        template: {
-          spec: {
-            containers: [
-              baseContainer + containerArgs + resources,
-            ],
-            restartPolicy: "OnFailure",
-          },
-        },
-        replicaType: replica_type,
-      }
-    else {},
-
-  pytorchJobReplicaV1alpha2(replica_type, number, args, image, num_gpus=0)::
     local baseContainer = {
       image: image,
       name: "pytorch",
@@ -90,10 +58,7 @@ local num_masters = params.numMasters;
 local num_workers = params.numWorkers;
 local num_gpus = params.numGpus;
 
-local replicaSpec = if jobVersion == "v1alpha1" then
-  util.pytorchJobReplica
-else
-  util.pytorchJobReplicaV1alpha2;
+local replicaSpec = util.pytorchJobReplica;
 
 local workerSpec = if num_gpus > 0 then
   replicaSpec("WORKER", num_workers, args, image, num_gpus)
@@ -102,19 +67,8 @@ else
 
 local masterSpec = replicaSpec("MASTER", num_masters, args, image);
 
-local job = if jobVersion == "v1alpha1" then {
-  apiVersion: "kubeflow.org/v1alpha1",
-  kind: "PyTorchJob",
-  metadata: {
-    name: name,
-    namespace: namespace,
-  },
-  spec: {
-    replicaSpecs: [masterSpec, workerSpec],
-  },
-
-} else {
-  apiVersion: "kubeflow.org/v1alpha2",
+local job = if jobVersion == "v1beta1" || jobVersion == "v1alpha2" then {
+  apiVersion: "kubeflow.org/" + jobVersion,
   kind: "PyTorchJob",
   metadata: {
     name: name,
@@ -126,6 +80,6 @@ local job = if jobVersion == "v1alpha1" then {
       Worker: workerSpec,
     },
   },
-};
+} else {};
 
 std.prune(k.core.v1.list.new([job]))
diff --git a/kubeflow/pytorch-job/prototypes/pytorch-operator.jsonnet b/kubeflow/pytorch-job/prototypes/pytorch-operator.jsonnet
@@ -5,9 +5,9 @@
 // @param name string Name to give to each of the components
 // @optionalParam disks string null Comma separated list of Google persistent disks to attach to jupyter environments.
 // @optionalParam cloud string null String identifying the cloud to customize the deployment for.
-// @optionalParam pytorchJobImage string gcr.io/kubeflow-images-public/pytorch-operator:v0.3.0 The image for the PyTorchJob controller
+// @optionalParam pytorchJobImage string gcr.io/kubeflow-images-public/pytorch-operator:v0.3.0-16-g8d71ff3 The image for the PyTorchJob controller
 // @optionalParam pytorchDefaultImage string null The default image to use for pytorch
-// @optionalParam pytorchJobVersion string v1alpha2 which version of the PyTorchJob operator to use
+// @optionalParam pytorchJobVersion string v1beta1 which version of the PyTorchJob operator to use
 // @optionalParam deploymentScope string cluster The scope at which pytorch-operator should be deployed - valid values are cluster, namespace.
 // @optionalParam deploymentNamespace string null The namespace to which pytorch-operator should be scoped. If deploymentScope is set to cluster, this is ignored.
 

diff --git a/kubeflow/pytorch-job/pytorch-operator.libsonnet b/kubeflow/pytorch-job/pytorch-operator.libsonnet
@@ -11,35 +11,14 @@
                          $.parts(params, env).crdV1beta1,
                          $.parts(params, env).pytorchJobDeployV1beta1(params.pytorchJobImage, params.deploymentScope, params.deploymentNamespace),
                        ]
-                     else if params.pytorchJobVersion == "v1alpha2" then
+                     else
                        [
                          $.parts(params, env).crdV1alpha2,
                          $.parts(params, env).pytorchJobDeployV1alpha2(params.pytorchJobImage, params.deploymentScope, params.deploymentNamespace),
-                       ]
-                     else
-                       [
-                         $.parts(params, env).crd,
-                         $.parts(params, env).pytorchJobDeploy(params.pytorchJobImage),
                        ],
 
   parts(params, env):: {
     local namespace = env.namespace,
-    crd: {
-      apiVersion: "apiextensions.k8s.io/v1beta1",
-      kind: "CustomResourceDefinition",
-      metadata: {
-        name: "pytorchjobs.kubeflow.org",
-      },
-      spec: {
-        group: "kubeflow.org",
-        version: "v1alpha1",
-        names: {
-          kind: "PyTorchJob",
-          singular: "pytorchjob",
-          plural: "pytorchjobs",
-        },
-      },
-    },
 
     crdV1alpha2: {
       apiVersion: "apiextensions.k8s.io/v1beta1",
@@ -137,72 +116,6 @@
       },
     },
 
-    pytorchJobDeploy(image): {
-      apiVersion: "extensions/v1beta1",
-      kind: "Deployment",
-      metadata: {
-        name: "pytorch-operator",
-        namespace: namespace,
-      },
-      spec: {
-        replicas: 1,
-        template: {
-          metadata: {
-            labels: {
-              name: "pytorch-operator",
-            },
-          },
-          spec: {
-            containers: [
-              {
-                command: [
-                  "/pytorch-operator",
-                  "--controller-config-file=/etc/config/controller_config_file.yaml",
-                  "--alsologtostderr",
-                  "-v=1",
-                ],
-                env: [
-                  {
-                    name: "MY_POD_NAMESPACE",
-                    valueFrom: {
-                      fieldRef: {
-                        fieldPath: "metadata.namespace",
-                      },
-                    },
-                  },
-                  {
-                    name: "MY_POD_NAME",
-                    valueFrom: {
-                      fieldRef: {
-                        fieldPath: "metadata.name",
-                      },
-                    },
-                  },
-                ],
-                image: image,
-                name: "pytorch-operator",
-                volumeMounts: [
-                  {
-                    mountPath: "/etc/config",
-                    name: "config-volume",
-                  },
-                ],
-              },
-            ],
-            serviceAccountName: "pytorch-operator",
-            volumes: [
-              {
-                configMap: {
-                  name: "pytorch-operator-config",
-                },
-                name: "config-volume",
-              },
-            ],
-          },
-        },
-      },
-    },  // pytorchJobDeploy
-
     pytorchJobDeployV1alpha2(image, deploymentScope, deploymentNamespace): {
       apiVersion: "extensions/v1beta1",
       kind: "Deployment",