Add pytorch framework specification
@@ -154,6 +154,63 @@ class HorovodSpecification(object):
return total_resources.to_dict()
+class PytorchSpecification(object):
+
+ @staticmethod
+ def get_cluster_def(cluster, pytorch_config):
+ is_distributed = False
+ if not pytorch_config:
+ return cluster, is_distributed
+
+ cluster[TaskType.WORKER] = pytorch_config.n_workers
+ if pytorch_config.n_workers != 0:
+ is_distributed = True
+
+ return cluster, is_distributed
+
+ @staticmethod
+ def get_worker_configs(environment, cluster, is_distributed):
+ if environment is None or environment.pytorch is None:
+ return {}
+ return get_task_configs(cluster=cluster,
+ is_distributed=is_distributed,
+ configs=environment.pytorch.worker_configs,
+ default_config=environment.pytorch.default_worker_config,
+ task_type=TaskType.WORKER)
+
+ @staticmethod
+ def get_worker_resources(environment, cluster, is_distributed):
+ if environment is None or environment.pytorch is None:
+ return None
+ return get_task_job_resources(
+ cluster=cluster,
+ is_distributed=is_distributed,
+ resources=environment.pytorch.worker_resources,
+ default_resources=environment.pytorch.default_worker_resources,
+ task_type=TaskType.WORKER)
+
+ @classmethod
+ def get_total_resources(cls, master_resources, environment, cluster, is_distributed):
+ worker_resources = cls.get_worker_resources(
+ environment=environment,
+ cluster=cluster,
+ is_distributed=is_distributed,
+ )
+
+ if not any([master_resources, worker_resources]):
+ return None
+
+ total_resources = PodResourcesConfig()
+
+ if master_resources:
+ total_resources += master_resources
+
+ for w_resources in six.itervalues(worker_resources or {}):
+ total_resources += w_resources
+
+ return total_resources.to_dict()
+
+
class MXNetSpecification(object):
@staticmethod
