def get_mpi_job_custom_resource(
resource_name: str,
namespace: str,
launcher: Optional[ReplicaSpec],
worker: Optional[ReplicaSpec],
termination: V1Termination,
collect_logs: bool,
sync_statuses: bool,
notifications: List[V1Notification],
clean_pod_policy: Optional[str],
slots_per_worker: Optional[int],
labels: Dict[str, str],
annotations: Dict[str, str],
) -> Dict:
template_spec = {}
get_kf_replicas_template(
replica_name="Launcher",
replica=launcher,
namespace=namespace,
resource_name=resource_name,
labels=labels,
annotations=annotations,
template_spec=template_spec,
)
get_kf_replicas_template(
replica_name="Worker",
replica=worker,
namespace=namespace,
resource_name=resource_name,
labels=labels,
annotations=annotations,
template_spec=template_spec,
)
template_spec = {"replicaSpecs": template_spec}
template_spec = set_clean_pod_policy(template_spec=template_spec,
clean_pod_policy=clean_pod_policy)
template_spec = set_slots_per_worker(template_spec=template_spec,
slots_per_worker=slots_per_worker)
custom_object = {"mpiJobSpec": template_spec}
custom_object = set_termination(custom_object=custom_object,
termination=termination)
custom_object = set_collect_logs(custom_object=custom_object,
collect_logs=collect_logs)
custom_object = set_sync_statuses(custom_object=custom_object,
sync_statuses=sync_statuses)
custom_object = set_notify(custom_object=custom_object,
notifications=notifications)
return get_operation_custom_object(
namespace=namespace,
resource_name=resource_name,
labels=labels,
annotations=annotations,
custom_object=custom_object,
)
def get_pytorch_job_custom_resource(
resource_name: str,
namespace: str,
master: Optional[ReplicaSpec],
worker: Optional[ReplicaSpec],
termination: V1Termination,
collect_logs: bool,
sync_statuses: bool,
notifications: List[V1Notification],
clean_pod_policy: Optional[str],
labels: Dict[str, str],
) -> Dict:
template_spec = {}
get_kf_replicas_template(
replica_name="Master",
replica=master,
namespace=namespace,
resource_name=resource_name,
labels=labels,
template_spec=template_spec,
)
get_kf_replicas_template(
replica_name="Worker",
replica=worker,
namespace=namespace,
resource_name=resource_name,
labels=labels,
template_spec=template_spec,
)
template_spec = set_clean_pod_policy(
template_spec=template_spec, clean_pod_policy=clean_pod_policy
)
custom_object = {"pytorchJobSpec": template_spec}
custom_object = set_termination(
custom_object=custom_object, termination=termination
)
custom_object = set_collect_logs(
custom_object=custom_object, collect_logs=collect_logs
)
custom_object = set_sync_statuses(
custom_object=custom_object, sync_statuses=sync_statuses
)
custom_object = set_notify(custom_object=custom_object, notifications=notifications)
return get_operation_custom_object(
namespace=namespace,
resource_name=resource_name,
labels=labels,
custom_object=custom_object,
)
def get_service_custom_resource(
resource_name: str,
namespace: str,
main_container: k8s_schemas.V1Container,
sidecar_containers: Optional[List[k8s_schemas.V1Container]],
init_containers: Optional[List[k8s_schemas.V1Container]],
volumes: List[k8s_schemas.V1Volume],
termination: V1Termination,
collect_logs: bool,
sync_statuses: bool,
notifications: List[V1Notification],
environment: V1Environment,
ports: List[int],
labels: Dict[str, str],
annotations: Dict[str, str],
) -> Dict:
metadata, pod_spec = get_pod_spec(
namespace=namespace,
main_container=main_container,
sidecar_containers=sidecar_containers,
init_containers=init_containers,
resource_name=resource_name,
volumes=volumes,
environment=environment,
labels=labels,
)
template_spec = {
"template": get_pod_template_spec(metadata=metadata, pod_spec=pod_spec)
}
if ports:
template_spec["ports"] = ports
custom_object = {"serviceSpec": template_spec}
custom_object = set_termination(
custom_object=custom_object, termination=termination
)
custom_object = set_collect_logs(
custom_object=custom_object, collect_logs=collect_logs
)
custom_object = set_sync_statuses(
custom_object=custom_object, sync_statuses=sync_statuses
)
custom_object = set_notify(custom_object=custom_object, notifications=notifications)
return get_operation_custom_object(
namespace=namespace,
resource_name=resource_name,
labels=labels,
annotations=annotations,
custom_object=custom_object,
)
def get_tf_job_custom_resource(
resource_name: str,
namespace: str,
chief: Optional[ReplicaSpec],
worker: Optional[ReplicaSpec],
ps: Optional[ReplicaSpec],
evaluator: Optional[ReplicaSpec],
termination: V1Termination,
collect_logs: bool,
sync_statuses: bool,
notifications: List[V1Notification],
clean_pod_policy: Optional[str],
labels: Dict[str, str],
annotations: Dict[str, str],
) -> Dict:
template_spec = {}
get_kf_replicas_template(
replica_name="Chief",
replica=chief,
namespace=namespace,
resource_name=resource_name,
labels=labels,
template_spec=template_spec,
)
get_kf_replicas_template(
replica_name="Worker",
replica=worker,
namespace=namespace,
resource_name=resource_name,
labels=labels,
template_spec=template_spec,
)
get_kf_replicas_template(
replica_name="PS",
replica=ps,
namespace=namespace,
resource_name=resource_name,
labels=labels,
template_spec=template_spec,
)
get_kf_replicas_template(
replica_name="Evaluator",
replica=evaluator,
namespace=namespace,
resource_name=resource_name,
labels=labels,
template_spec=template_spec,
)
template_spec = {"replicaSpecs": template_spec}
template_spec = set_clean_pod_policy(
template_spec=template_spec, clean_pod_policy=clean_pod_policy
)
custom_object = {"tfJobSpec": template_spec}
custom_object = set_termination(
custom_object=custom_object, termination=termination
)
custom_object = set_collect_logs(
custom_object=custom_object, collect_logs=collect_logs
)
custom_object = set_sync_statuses(
custom_object=custom_object, sync_statuses=sync_statuses
)
custom_object = set_notify(custom_object=custom_object, notifications=notifications)
return get_operation_custom_object(
namespace=namespace,
resource_name=resource_name,
labels=labels,
annotations=annotations,
custom_object=custom_object,
)
