fixed prototype to set termination policy to worker 0 by default when…

… no master exists. (#444)
kubeflow · Mar 19, 2018 · 5dca627 · 5dca627
1 parent 6073021
commit 5dca627
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 3 deletions.
diff --git a/kubeflow/tf-job/prototypes/tf-job.jsonnet b/kubeflow/tf-job/prototypes/tf-job.jsonnet
@@ -37,6 +37,11 @@ local numPs = import "param://num_ps";
 local numWorkers = import "param://num_workers";
 local numGpus = import "param://num_gpus";
 
+local terminationPolicy = if numMasters == 1 then
+  tfJob.parts.tfJobTerminationPolicy("MASTER", 0)
+else
+  tfJob.parts.tfJobTerminationPolicy("WORKER", 0);
+
 local workerSpec = if numGpus > 0 then
   tfJob.parts.tfJobReplica("WORKER", numWorkers, args, imageGpu, numGpus)
 else
@@ -46,6 +51,7 @@ std.prune(k.core.v1.list.new([
   tfJob.parts.tfJob(name, namespace, [
     tfJob.parts.tfJobReplica("MASTER", numMasters, args, image),
     workerSpec,
-    tfJob.parts.tfJobReplica("PS", numPs, args, image),
-  ]),
+    tfJob.parts.tfJobReplica("PS", numPs, args, image),],
+    terminationPolicy
+  ),
 ]))
diff --git a/kubeflow/tf-job/tf-job.libsonnet b/kubeflow/tf-job/tf-job.libsonnet
@@ -34,7 +34,14 @@ local k = import "k.libsonnet";
         }
       else {},
 
-    tfJob(name, namespace, replicas):: {
+    tfJobTerminationPolicy(replicaName, replicaIndex):: {
+      chief: {
+          replicaName: replicaName,
+          replicaIndex: replicaIndex,
+      },
+    },
+
+    tfJob(name, namespace, replicas, tp):: {
       apiVersion: "kubeflow.org/v1alpha1",
       kind: "TFJob",
       metadata: {
@@ -43,6 +50,7 @@ local k = import "k.libsonnet";
       },
       spec: {
         replicaSpecs: replicas,
+        terminationPolicy: tp,
       },
     },
   },