def test_get_job_custom_resource(self):
termination = V1Termination(max_retries=5, ttl=10, timeout=10)
environment = V1Environment(
labels={"foo": "bar"},
annotations={"foo": "bar"},
node_selector={"foo": "bar"},
node_name="foo",
restart_policy="never",
)
notifications = [
V1Notification(connections=["test"], trigger=V1Statuses.DONE)
]
chief, chief_replica_template = self.get_replica(environment)
worker, worker_replica_template = self.get_replica(environment)
template_spec = {
"cleanPodPolicy": "Running",
"replicaSpecs": {
"Chief": chief_replica_template,
"Worker": worker_replica_template,
},
}
custom_object = {
"tfJobSpec": template_spec,
"termination": {
"backoffLimit": termination.max_retries,
"activeDeadlineSeconds": termination.timeout,
"ttlSecondsAfterFinished": termination.ttl,
},
"collectLogs": True,
"syncStatuses": True,
"notifications": [n.to_operator() for n in notifications],
}
expected_crd = get_custom_object(
namespace="default",
resource_name="foo",
kind="Operation",
api_version="core.polyaxon.com/v1",
labels={"foo": "bar"},
annotations={"foo": "bar"},
custom_object=custom_object,
)
crd = get_tf_job_custom_resource(
namespace="default",
resource_name="foo",
chief=chief,
worker=worker,
ps=None,
evaluator=None,
clean_pod_policy="Running",
termination=termination,
collect_logs=True,
sync_statuses=True,
notifications=notifications,
labels=environment.labels,
annotations={"foo": "bar"},
)
assert crd == expected_crd
def get_resource(
self,
compiled_operation: V1CompiledOperation,
artifacts_store: V1ConnectionType,
connection_by_names: Dict[str, V1ConnectionType],
secrets: Optional[Iterable[V1K8sResourceType]],
config_maps: Optional[Iterable[V1K8sResourceType]],
default_sa: str = None,
default_auth: bool = False,
) -> Dict:
job = compiled_operation.run # type: V1TFJob
def _get_replica(
replica: Optional[V1KFReplica]) -> Optional[ReplicaSpec]:
if not replica:
return None
return self.get_replica_resource(
plugins=plugins,
contexts=contexts,
environment=replica.environment,
volumes=replica.volumes or [],
init=replica.init or [],
sidecars=replica.sidecars or [],
container=replica.container,
artifacts_store=artifacts_store,
connections=replica.connections or [],
connection_by_names=connection_by_names,
secrets=secrets,
config_maps=config_maps,
default_sa=default_sa,
num_replicas=replica.replicas,
)
plugins = compiled_operation.plugins or V1Plugins()
contexts = PluginsContextsSpec.from_config(plugins,
default_auth=default_auth)
chief = _get_replica(job.chief)
worker = _get_replica(job.worker)
ps = _get_replica(job.ps)
evaluator = _get_replica(job.evaluator)
labels = self.get_labels(version=pkg.VERSION, labels={})
return get_tf_job_custom_resource(
namespace=self.namespace,
resource_name=self.resource_name,
chief=chief,
worker=worker,
ps=ps,
evaluator=evaluator,
termination=compiled_operation.termination,
collect_logs=contexts.collect_logs,
clean_pod_policy=job.clean_pod_policy,
sync_statuses=contexts.sync_statuses,
notifications=plugins.notifications,
labels=labels,
annotations=self.annotations,
)
def test_get_tf_job_custom_resource_with_no_workers(self):
termination = V1Termination(max_retries=5, ttl=10, timeout=10)
environment = V1Environment(
labels={"foo": "bar"},
annotations={"foo": "bar"},
node_selector={"foo": "bar"},
node_name="foo",
restart_policy="never",
)
custom_object = {
"tfJobSpec": {
"cleanPodPolicy": "All",
"replicaSpecs": {}
},
"termination": {
"backoffLimit": termination.max_retries,
"activeDeadlineSeconds": termination.timeout,
"ttlSecondsAfterFinished": termination.ttl,
},
"collectLogs": False,
"syncStatuses": False,
"notifications": [],
}
expected_crd = get_custom_object(
namespace="default",
resource_name="foo",
kind="Operation",
api_version="core.polyaxon.com/v1",
labels={"foo": "bar"},
annotations={"foo": "long-foo-bar" * 300},
custom_object=custom_object,
)
crd = get_tf_job_custom_resource(
namespace="default",
resource_name="foo",
chief=None,
worker=None,
ps=None,
evaluator=None,
clean_pod_policy=None,
termination=termination,
collect_logs=False,
sync_statuses=False,
notifications=None,
labels=environment.labels,
annotations={"foo": "long-foo-bar" * 300},
)
assert crd == expected_crd
