avoid launching both chief and worker jobs when running a non distrib…

…uted job (#50)
kubeflow · Jan 9, 2019 · 57a59d6 · 57a59d6
1 parent 6c1f5bb
commit 57a59d6
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 5 deletions.
diff --git a/fairing/training/kubeflow/decorators.py b/fairing/training/kubeflow/decorators.py
@@ -25,6 +25,7 @@ def __init__(self, worker_count=0, ps_count=0, namespace=None):
         super(DistributedTraining, self).__init__(namespace)
         self.distribution = {
             'Worker': worker_count,
-            'PS': ps_count
+            'PS': ps_count,
+            'Chief': 1
         }
 
diff --git a/fairing/training/kubeflow/deployment.py b/fairing/training/kubeflow/deployment.py
@@ -20,17 +20,18 @@ def generate_job(self, pod_template_spec):
         worker_replica_spec['template'] = pod_template_spec
 
         ps_replica_spec = {}
-        ps_replica_spec['replicas'] = self.distribution['PS']
+        ps_replica_spec['replicas'] = self.distribution.get('PS', 0)
         ps_replica_spec['template'] = pod_template_spec
 
         chief_replica_spec = {}
-        chief_replica_spec['replicas'] = 1
+        chief_replica_spec['replicas'] = self.distribution.get('Chief', 0)
         chief_replica_spec['template'] = pod_template_spec
 
         spec = {} 
         spec['tfReplicaSpecs'] = {}
-        spec['tfReplicaSpecs']['Chief'] = chief_replica_spec
         spec['tfReplicaSpecs']['Worker'] = worker_replica_spec
+        if chief_replica_spec['replicas'] > 0:
+            spec['tfReplicaSpecs']['Chief'] = chief_replica_spec
         if ps_replica_spec['replicas'] > 0:
             spec['tfReplicaSpecs']['PS'] = ps_replica_spec
 

diff --git a/setup.py b/setup.py
@@ -29,7 +29,7 @@
         'notebook==5.6.0',
         'jupyter==1.0.0',
         'numpy==1.15.0',
-        'kubernetes==6.0.0',
+        'kubernetes==8.0.1',
         'future==0.17.1',
         'six==1.11.0',
         'httplib2==0.12.0',