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=V1NotificationTrigger.DONE) ] master, master_replica_template = self.get_replica(environment) worker, worker_replica_template = self.get_replica(environment) template_spec = { "cleanPodPolicy": "Running", "replicaSpecs": { "Master": master_replica_template, "Worker": worker_replica_template, }, } custom_object = { "pytorchJobSpec": 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_pytorch_job_custom_resource( namespace="default", resource_name="foo", master=master, worker=worker, clean_pod_policy="Running", termination=termination, labels=environment.labels, annotations={"foo": "bar"}, notifications=notifications, collect_logs=True, sync_statuses=True, ) 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: V1PytorchJob 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) master = _get_replica(job.master) worker = _get_replica(job.worker) labels = self.get_labels(version=pkg.VERSION, labels={}) return get_pytorch_job_custom_resource( namespace=self.namespace, resource_name=self.resource_name, master=master, worker=worker, termination=compiled_operation.termination, clean_pod_policy=job.clean_pod_policy, collect_logs=contexts.collect_logs, sync_statuses=contexts.sync_statuses, notifications=plugins.notifications, labels=labels, annotations=self.annotations, )
def test_get_pytorch_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 = { "pytorchJobSpec": { "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"}, custom_object=custom_object, annotations={"foo": "long-foo-bar" * 300}, ) crd = get_pytorch_job_custom_resource( namespace="default", resource_name="foo", master=None, worker=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