def affinities(self): cluster, is_distributed, = self.spec.cluster_def worker_affinities = HorovodSpecification.get_worker_affinities( environment=self.spec.environment, cluster=cluster, is_distributed=is_distributed) return { TaskType.MASTER: { 0: self.spec.master_affinity }, TaskType.WORKER: worker_affinities, }
def handle_horovod_experiment(experiment, spawner, response): # Get the number of jobs this experiment started master = response[TaskType.MASTER] job_uuid = master['pod']['metadata']['labels']['job_uuid'] job_uuid = uuid.UUID(job_uuid) create_job(job_uuid=job_uuid, experiment=experiment, definition=get_job_definition(master), resources=spawner.spec.master_resources, node_selector=spawner.spec.master_node_selector, affinity=spawner.spec.master_affinity, tolerations=spawner.spec.master_tolerations) cluster, is_distributed = spawner.spec.cluster_def worker_resources = HorovodSpecification.get_worker_resources( environment=spawner.spec.environment, cluster=cluster, is_distributed=is_distributed ) worker_node_selectors = HorovodSpecification.get_worker_node_selectors( environment=spawner.spec.environment, cluster=cluster, is_distributed=is_distributed ) worker_affinities = HorovodSpecification.get_worker_affinities( environment=spawner.spec.environment, cluster=cluster, is_distributed=is_distributed ) worker_tolerations = HorovodSpecification.get_worker_tolerations( environment=spawner.spec.environment, cluster=cluster, is_distributed=is_distributed ) for i, worker in enumerate(response[TaskType.WORKER]): job_uuid = worker['pod']['metadata']['labels']['job_uuid'] job_uuid = uuid.UUID(job_uuid) create_job(job_uuid=job_uuid, experiment=experiment, definition=get_job_definition(worker), role=TaskType.WORKER, resources=worker_resources.get(i), node_selector=worker_node_selectors.get(i), affinity=worker_affinities.get(i), tolerations=worker_tolerations.get(i))