polyaxon 2.6.1__py3-none-any.whl → 2.7.0rc0__py3-none-any.whl

polyaxon/_deploy/schemas/deployment.py

@@ -183,6 +183,7 @@ def validate_deployment_chart(
 
 
 class DeploymentConfig(BaseSchemaModel):
+    _SWAGGER_FIELDS_LISTS = ["tolerations", "celeryTolerations"]
     _SWAGGER_FIELDS = [
         "tolerations",
         "affinity",

polyaxon/_flow/component/component.py

@@ -396,3 +396,7 @@ class V1Component(
 
     def get_name(self):
         return self.name.split(":")[0] if self.name else None
+
+    def get_replica_types(self):
+        if self.is_distributed_run:
+            return self.run.get_replica_types()

polyaxon/_flow/operations/compiled_operation.py

@@ -135,3 +135,7 @@ class V1CompiledOperation(BaseOp, RunMixin):
             raise PolyaxonSchemaError(
                 "Operations with dag runtime do not support the `build` section."
             )
+
+    def get_replica_types(self):
+        if self.is_distributed_run:
+            return self.run.get_replica_types()

polyaxon/_flow/operations/operation.py

@@ -648,8 +648,11 @@ class V1Operation(BaseOp, TemplateMixinConfig):
             return result
 
         kind = config.component.run.kind
-        value = validate_run_patch(value, kind)
-        current_value = validate_run_patch(current_value, kind)
+        replica_types = config.component.get_replica_types()
+        value = validate_run_patch(value, kind, replica_types=replica_types)
+        current_value = validate_run_patch(
+            current_value, kind, replica_types=replica_types
+        )
         run_patch = current_value.patch(value, strategy)
         run_patch = run_patch.to_dict()
         run_patch.pop("kind")

polyaxon/_flow/run/__init__.py

@@ -52,6 +52,9 @@ V1Runtime = Annotated[
 
 
 class RunMixin:
+    def get_replica_types(self):
+        raise NotImplementedError
+
     def get_run_kind(self):
         raise NotImplementedError
 

polyaxon/_flow/run/dask/dask.py

@@ -175,3 +175,13 @@ class V1DaskJob(BaseRun, DestinationImageMixin):
         if self.scheduler:
             init += self.scheduler.get_all_init()
         return init
+
+    def get_replica_types(self):
+        types = []
+        if self.job:
+            types.append("job")
+        if self.worker:
+            types.append("worker")
+        if self.scheduler:
+            types.append("scheduler")
+        return types

polyaxon/_flow/run/kubeflow/mpi_job.py

@@ -185,3 +185,11 @@ class V1MPIJob(BaseRun, DestinationImageMixin):
         if self.worker:
             init += self.worker.get_all_init()
         return init
+
+    def get_replica_types(self):
+        types = []
+        if self.launcher:
+            types.append("launcher")
+        if self.worker:
+            types.append("worker")
+        return types

polyaxon/_flow/run/kubeflow/mx_job.py

@@ -313,3 +313,19 @@ class V1MXJob(BaseRun, DestinationImageMixin):
         if self.tuner_server:
             init += self.tuner_server.get_all_init()
         return init
+
+    def get_replica_types(self):
+        types = []
+        if self.scheduler:
+            types.append("scheduler")
+        if self.server:
+            types.append("server")
+        if self.worker:
+            types.append("worker")
+        if self.tuner:
+            types.append("tuner")
+        if self.tuner_tracker:
+            types.append("tuner_tracker")
+        if self.tuner_server:
+            types.append("tuner_server")
+        return types

polyaxon/_flow/run/kubeflow/paddle_job.py

@@ -205,3 +205,11 @@ class V1PaddleJob(BaseRun, DestinationImageMixin):
         if self.worker:
             init += self.worker.get_all_init()
         return init
+
+    def get_replica_types(self):
+        types = []
+        if self.master:
+            types.append("master")
+        if self.worker:
+            types.append("worker")
+        return types

polyaxon/_flow/run/kubeflow/pytorch_job.py

@@ -218,3 +218,11 @@ class V1PytorchJob(BaseRun, DestinationImageMixin):
         if self.worker:
             init += self.worker.get_all_init()
         return init
+
+    def get_replica_types(self):
+        types = []
+        if self.master:
+            types.append("master")
+        if self.worker:
+            types.append("worker")
+        return types

polyaxon/_flow/run/kubeflow/tf_job.py

@@ -248,3 +248,15 @@ class V1TFJob(BaseRun, DestinationImageMixin):
         if self.evaluator:
             init += self.evaluator.get_all_init()
         return init
+
+    def get_replica_types(self):
+        types = []
+        if self.chief:
+            types.append("chief")
+        if self.ps:
+            types.append("ps")
+        if self.worker:
+            types.append("worker")
+        if self.evaluator:
+            types.append("evaluator")
+        return types

polyaxon/_flow/run/kubeflow/xgboost_job.py

@@ -130,63 +130,49 @@ class V1XGBoostJob(BaseRun, DestinationImageMixin):
     worker: Optional[Union[V1KFReplica, RefField]] = None
 
     def apply_image_destination(self, image: str):
-        if self.chief:
-            self.chief.container = self.chief.container or V1Container()
-            self.chief.container.image = image
-        if self.ps:
-            self.ps.container = self.ps.container or V1Container()
-            self.ps.container.image = image
+        if self.master:
+            self.master.container = self.master.container or V1Container()
+            self.master.container.image = image
         if self.worker:
             self.worker.container = self.worker.container or V1Container()
             self.worker.container.image = image
-        if self.evaluator:
-            self.evaluator.container = self.evaluator.container or V1Container()
-            self.evaluator.container.image = image
 
     def get_resources(self):
         resources = V1RunResources()
-        if self.chief:
-            resources += self.chief.get_resources()
-        if self.ps:
-            resources += self.ps.get_resources()
+        if self.master:
+            resources += self.master.get_resources()
         if self.worker:
             resources += self.worker.get_resources()
-        if self.evaluator:
-            resources += self.evaluator.get_resources()
         return resources
 
     def get_all_containers(self):
         containers = []
-        if self.chief:
-            containers += self.chief.get_all_containers()
-        if self.ps:
-            containers += self.ps.get_all_containers()
+        if self.master:
+            containers += self.master.get_all_containers()
         if self.worker:
             containers += self.worker.get_all_containers()
-        if self.evaluator:
-            containers += self.evaluator.get_all_containers()
         return containers
 
     def get_all_connections(self):
         connections = []
-        if self.chief:
-            connections += self.chief.get_all_connections()
-        if self.ps:
-            connections += self.ps.get_all_connections()
+        if self.master:
+            connections += self.master.get_all_connections()
         if self.worker:
             connections += self.worker.get_all_connections()
-        if self.evaluator:
-            connections += self.evaluator.get_all_connections()
         return connections
 
     def get_all_init(self):
         init = []
-        if self.chief:
-            init += self.chief.get_all_init()
-        if self.ps:
-            init += self.ps.get_all_init()
+        if self.master:
+            init += self.master.get_all_init()
         if self.worker:
             init += self.worker.get_all_init()
-        if self.evaluator:
-            init += self.evaluator.get_all_init()
         return init
+
+    def get_replica_types(self):
+        types = []
+        if self.master:
+            types.append(self.master.replicas)
+        if self.worker:
+            types.append(self.worker.replicas)
+        return types

polyaxon/_flow/run/patch.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Dict, List
 
 from clipped.compact.pydantic import ValidationError
 
@@ -21,7 +21,9 @@ from polyaxon._flow.run.tuner import V1TunerJob
 from polyaxon.exceptions import PolyaxonValidationError
 
 
-def validate_run_patch(run_patch: Dict, kind: V1RunKind):
+def validate_run_patch(
+    run_patch: Dict, kind: V1RunKind, replica_types: List[str] = None
+):
     if kind == V1RunKind.JOB:
         patch = V1Job.from_dict(run_patch)
     elif kind == V1RunKind.SERVICE:
@@ -37,37 +39,63 @@ def validate_run_patch(run_patch: Dict, kind: V1RunKind):
         try:
             patch = V1PytorchJob.from_dict(run_patch)
         except ValidationError:
-            patch = V1KFReplica.from_dict(run_patch)
+            if replica_types:
+                patch = V1PytorchJob.from_dict({k: run_patch for k in replica_types})
+            else:
+                patch = V1KFReplica.from_dict(run_patch)
     elif kind == V1RunKind.PADDLEJOB:
         try:
             patch = V1PaddleJob.from_dict(run_patch)
         except ValidationError:
-            patch = V1KFReplica.from_dict(run_patch)
+            if replica_types:
+                patch = V1PaddleJob.from_dict({k: run_patch for k in replica_types})
+            else:
+                patch = V1KFReplica.from_dict(run_patch)
     elif kind == V1RunKind.TFJOB:
         try:
             patch = V1TFJob.from_dict(run_patch)
         except ValidationError:
-            patch = V1KFReplica.from_dict(run_patch)
+            if replica_types:
+                patch = V1TFJob.from_dict({k: run_patch for k in replica_types})
+            else:
+                patch = V1KFReplica.from_dict(run_patch)
     elif kind == V1RunKind.MXJOB:
         try:
             patch = V1MXJob.from_dict(run_patch)
         except ValidationError:
-            patch = V1KFReplica.from_dict(run_patch)
+            if replica_types:
+                patch = V1MXJob.from_dict({k: run_patch for k in replica_types})
+            else:
+                patch = V1KFReplica.from_dict(run_patch)
     elif kind == V1RunKind.XGBJOB:
         try:
             patch = V1XGBoostJob.from_dict(run_patch)
         except ValidationError:
-            patch = V1KFReplica.from_dict(run_patch)
+            if replica_types:
+                patch = V1XGBoostJob.from_dict({k: run_patch for k in replica_types})
+            else:
+                patch = V1KFReplica.from_dict(run_patch)
     elif kind == V1RunKind.RAYJOB:
         try:
             patch = V1RayJob.from_dict(run_patch)
         except ValidationError:
-            patch = V1RayReplica.from_dict(run_patch)
+            if replica_types:
+                replicas = {}
+                if "head" in replica_types:
+                    replicas["head"] = run_patch
+                replica_types = [r for r in replica_types if r != "head"]
+                replicas["workers"] = {replica: run_patch for replica in replica_types}
+                patch = V1RayJob.from_dict(replicas)
+            else:
+                patch = V1RayReplica.from_dict(run_patch)
     elif kind == V1RunKind.DASKJOB:
         try:
             patch = V1DaskJob.from_dict(run_patch)
         except ValidationError:
-            patch = V1DaskReplica.from_dict(run_patch)
+            if replica_types:
+                patch = V1DaskJob.from_dict({k: run_patch for k in replica_types})
+            else:
+                patch = V1DaskReplica.from_dict(run_patch)
     elif kind == V1RunKind.NOTIFIER:
         patch = V1NotifierJob.from_dict(run_patch)
     elif kind == V1RunKind.TUNER:

polyaxon/_flow/run/ray/ray.py

@@ -136,6 +136,7 @@ class V1RayJob(BaseRun, DestinationImageMixin):
 
     _IDENTIFIER = V1RunKind.RAYJOB
     _CUSTOM_DUMP_FIELDS = {"head", "workers"}
+    _FIELDS_DICT_PATCH = ["workers"]
 
     kind: Literal[_IDENTIFIER] = _IDENTIFIER
     entrypoint: Optional[str] = None
@@ -196,3 +197,11 @@ class V1RayJob(BaseRun, DestinationImageMixin):
                 worker = self.workers[worker_name]
                 init += worker.get_all_init()
         return init
+
+    def get_replica_types(self):
+        types = []
+        if self.head:
+            types.append("head")
+        if self.workers:
+            types += list(self.workers.keys())
+        return types

polyaxon/_polyaxonfile/specs/compiled_operation.py

@@ -321,7 +321,11 @@ class CompiledOperationSpecification(BaseSpecification):
         preset = OperationSpecification.read(preset, is_preset=True)  # type: V1Operation
         if preset.run_patch:
             config.run = config.run.patch(
-                validate_run_patch(preset.run_patch, config.run.kind),
+                validate_run_patch(
+                    preset.run_patch,
+                    config.run.kind,
+                    replica_types=config.get_replica_types(),
+                ),
                 strategy=preset.patch_strategy,
             )
         patch_keys = {

polyaxon/_polyaxonfile/specs/operation.py

@@ -54,7 +54,11 @@ class OperationSpecification(BaseSpecification):
                 else config.patch_strategy
             )
             component.run = component.run.patch(
-                validate_run_patch(config.run_patch, component.run.kind),
+                validate_run_patch(
+                    config.run_patch,
+                    component.run.kind,
+                    replica_types=component.get_replica_types(),
+                ),
                 strategy=patch_strategy,
             )