Update scheduler to set correct k8s_replica for kubeflow
@@ -34,14 +34,19 @@ _logger = logging.getLogger('polyaxon.scheduler.experiment')
def create_job(job_uuid,
experiment,
role=None,
+ k8s_replica=None,
sequence=None,
resources=None,
node_selector=None,
affinity=None,
tolerations=None):
job = ExperimentJob(uuid=uuid.UUID(job_uuid), experiment=experiment, definition={})
+
+ k8s_replica = k8s_replica or role
if role:
job.role = role
+ if k8s_replica:
+ job.k8s_replica = k8s_replica
if sequence:
job.sequence = sequence
@@ -123,8 +128,12 @@ def get_spawner_class(backend, framework):
def create_tensorflow_experiment_jobs(experiment, spawner):
master_job_uuid = spawner.job_uuids[TaskType.MASTER][0]
+ k8s_replica = None
+ if experiment.backend == ExperimentBackend.KUBEFLOW:
+ k8s_replica = TaskType.CHIEF
create_job(job_uuid=master_job_uuid,
experiment=experiment,
+ k8s_replica=k8s_replica,
resources=spawner.spec.master_resources,
node_selector=spawner.spec.master_node_selector,
affinity=spawner.spec.master_affinity,
@@ -216,8 +225,12 @@ def handle_tensorflow_experiment(response):
def create_horovod_experiment_jobs(experiment, spawner):
master_job_uuid = spawner.job_uuids[TaskType.MASTER][0]
+ k8s_replica = None
+ if experiment.backend == ExperimentBackend.KUBEFLOW:
+ k8s_replica = TaskType.CHIEF
create_job(job_uuid=master_job_uuid,
experiment=experiment,
+ k8s_replica=k8s_replica,
resources=spawner.spec.master_resources,
node_selector=spawner.spec.master_node_selector,
affinity=spawner.spec.master_affinity,
@@ -272,8 +285,12 @@ def handle_horovod_experiment(response):
def create_pytorch_experiment_jobs(experiment, spawner):
master_job_uuid = spawner.job_uuids[TaskType.MASTER][0]
+ k8s_replica = None
+ if experiment.backend == ExperimentBackend.KUBEFLOW:
+ k8s_replica = TaskType.CHIEF
create_job(job_uuid=master_job_uuid,
experiment=experiment,
+ k8s_replica=k8s_replica,
resources=spawner.spec.master_resources,
node_selector=spawner.spec.master_node_selector,
affinity=spawner.spec.master_affinity,
@@ -329,8 +346,12 @@ def handle_pytorch_experiment(response):
def create_mxnet_experiment_jobs(experiment, spawner):
master_job_uuid = spawner.job_uuids[TaskType.MASTER][0]
+ k8s_replica = None
+ if experiment.backend == ExperimentBackend.KUBEFLOW:
+ k8s_replica = TaskType.CHIEF
create_job(job_uuid=master_job_uuid,
experiment=experiment,
+ k8s_replica=k8s_replica,
resources=spawner.spec.master_resources,
node_selector=spawner.spec.master_node_selector,
affinity=spawner.spec.master_affinity,
