polyaxon 2.1.0rc9__py3-none-any.whl → 2.6.0__py3-none-any.whl

polyaxon/_flow/operations/operation.py

@@ -2,7 +2,14 @@ from copy import copy
 from typing import Dict, Optional
 from typing_extensions import Literal
 
-from clipped.compact.pydantic import Field, StrictStr, root_validator, validator
+from clipped.compact.pydantic import (
+    Field,
+    StrictStr,
+    field_validator,
+    model_validator,
+    validation_after,
+    validation_before,
+)
 from clipped.config.patch_strategy import PatchStrategy
 from clipped.config.schema import skip_partial, to_partial
 
@@ -508,40 +515,41 @@ class V1Operation(BaseOp, TemplateMixinConfig):
         "run_patch",
         "patch_strategy",
     ]
+    _CUSTOM_DUMP_FIELDS = {"component"}
 
     kind: Literal[_IDENTIFIER] = _IDENTIFIER
-    params: Optional[Dict[StrictStr, V1Param]]
-    hub_ref: Optional[StrictStr] = Field(alias="hubRef")
-    dag_ref: Optional[StrictStr] = Field(alias="dagRef")
-    url_ref: Optional[StrictStr] = Field(alias="urlRef")
-    path_ref: Optional[StrictStr] = Field(alias="pathRef")
-    component: Optional[V1Component]
-    patch_strategy: Optional[PatchStrategy] = Field(alias="patchStrategy")
-    is_preset: Optional[bool] = Field(alias="isPreset")
-    run_patch: Optional[Dict] = Field(alias="runPatch")
-    template: Optional[V1Template]
-
-    @root_validator
+    params: Optional[Dict[StrictStr, V1Param]] = None
+    hub_ref: Optional[StrictStr] = Field(alias="hubRef", default=None)
+    dag_ref: Optional[StrictStr] = Field(alias="dagRef", default=None)
+    url_ref: Optional[StrictStr] = Field(alias="urlRef", default=None)
+    path_ref: Optional[StrictStr] = Field(alias="pathRef", default=None)
+    component: Optional[V1Component] = None
+    patch_strategy: Optional[PatchStrategy] = Field(alias="patchStrategy", default=None)
+    is_preset: Optional[bool] = Field(alias="isPreset", default=None)
+    run_patch: Optional[Dict] = Field(alias="runPatch", default=None)
+    template: Optional[V1Template] = None
+
+    @model_validator(**validation_after)
     @skip_partial
     def validate_reference(cls, values):
         if not values:
             return values
-        if values.get("is_preset"):
+        if cls.get_value_for_key("is_preset", values):
             return values
         count = 0
-        hub_ref = values.get("hub_ref")
+        hub_ref = cls.get_value_for_key("hub_ref", values)
         if hub_ref:
             count += 1
-        dag_ref = values.get("dag_ref")
+        dag_ref = cls.get_value_for_key("dag_ref", values)
         if dag_ref:
             count += 1
-        url_ref = values.get("url_ref")
+        url_ref = cls.get_value_for_key("url_ref", values)
         if url_ref:
             count += 1
-        path_ref = values.get("path_ref")
+        path_ref = cls.get_value_for_key("path_ref", values)
         if path_ref:
             count += 1
-        component = values.get("component")
+        component = cls.get_value_for_key("component", values)
         if component and count == 0:
             count += 1
 
@@ -552,11 +560,11 @@ class V1Operation(BaseOp, TemplateMixinConfig):
             )
         return values
 
-    @validator("run_patch")
+    @field_validator("run_patch")
     @skip_partial
     def validate_run_patch(cls, run_patch, values):
-        component = values.get("component")
-        if values.get("is_preset"):
+        component = cls.get_value_for_key("component", values)
+        if cls.get_value_for_key("is_preset", values):
             return run_patch
         if not component or not run_patch:
             return run_patch
@@ -662,14 +670,16 @@ class V1Operation(BaseOp, TemplateMixinConfig):
             for k, v in contexts.items():
                 params[k] = V1Param(value=v, context_only=True)
 
-        return cls.construct(
-            run_patch=run_patch,
-            hub_ref=hook.hub_ref,
-            presets=hook.presets,
-            queue=hook.queue,
-            namespace=hook.namespace,
-            params=params,
-        )
+        content = {"run_patch": run_patch, "params": params}
+        if hook.hub_ref:
+            content["hub_ref"] = hook.hub_ref
+        if hook.presets:
+            content["presets"] = hook.presets
+        if hook.queue:
+            content["queue"] = hook.queue
+        if hook.namespace:
+            content["namespace"] = hook.namespace
+        return cls.construct(**content)
 
     @classmethod
     def from_build(cls, build: V1Build, contexts: Optional[Dict] = None):
@@ -685,17 +695,23 @@ class V1Operation(BaseOp, TemplateMixinConfig):
         if not destination.connection or build.connection:
             destination.connection = build.connection
         params["destination"] = destination
-
-        return cls.construct(
-            run_patch=build.run_patch,
-            patch_strategy=build.patch_strategy,
-            hub_ref=build.hub_ref,
-            presets=build.presets,
-            queue=build.queue,
-            namespace=build.namespace,
-            cache=build.cache,
-            params=params,
-        )
+        content = {
+            "run_patch": build.run_patch,
+            "patch_strategy": build.patch_strategy,
+            "params": params,
+        }
+        if build.hub_ref:
+            content["hub_ref"] = build.hub_ref
+        if build.presets:
+            content["presets"] = build.presets
+        if build.queue:
+            content["queue"] = build.queue
+        if build.namespace:
+            content["namespace"] = build.namespace
+        if build.cache:
+            content["cache"] = build.cache
+
+        return cls.construct(**content)
 
 
 PartialV1Operation = to_partial(V1Operation)

polyaxon/_flow/optimization/__init__.py

@@ -10,7 +10,7 @@ class V1OptimizationMetric(BaseSchemaModel):
     _IDENTIFIER = "optimization_metric"
 
     name: StrictStr
-    optimization: Optional[V1Optimization]
+    optimization: Optional[V1Optimization] = None
 
     def get_for_sort(self):
         if self.optimization == V1Optimization.MINIMIZE:
@@ -22,7 +22,7 @@ class V1OptimizationResource(BaseSchemaModel):
     _IDENTIFIER = "optimization_resource"
 
     name: StrictStr
-    type: Optional[V1ResourceType]
+    type: Optional[V1ResourceType] = None
 
     def cast_value(self, value):
         if V1ResourceType.is_int(self.type):

polyaxon/_flow/params/params.py

@@ -2,7 +2,7 @@ from collections import namedtuple
 from collections.abc import Mapping
 from typing import Any, Dict, Optional
 
-from clipped.compact.pydantic import Field, StrictStr, validator
+from clipped.compact.pydantic import Field, StrictStr, field_validator
 from clipped.config.schema import skip_partial
 from clipped.utils.lists import to_list
 from clipped.utils.strings import to_string
@@ -18,7 +18,7 @@ from polyaxon.exceptions import PolyaxonValidationError
 
 def validate_param_value(value, ref):
     if ref and not isinstance(value, str):
-        raise TypeError(
+        raise ValueError(
             "Value must be of type string when a ref is provided, received `{}` instead.".format(
                 value
             )
@@ -387,16 +387,17 @@ class V1Param(BaseSchemaModel, ctx_refs.RefMixin, ParamValueMixin):
     _IDENTIFIER = "param"
 
     value: Any
-    ref: Optional[StrictStr]
-    context_only: Optional[bool] = Field(alias="contextOnly")
-    connection: Optional[StrictStr]
-    to_init: Optional[bool] = Field(alias="toInit")
-    to_env: Optional[StrictStr] = Field(alias="toEnv")
+    ref: Optional[StrictStr] = None
+    context_only: Optional[bool] = Field(alias="contextOnly", default=False)
+    connection: Optional[StrictStr] = None
+    to_init: Optional[bool] = Field(alias="toInit", default=None)
+    to_env: Optional[StrictStr] = Field(alias="toEnv", default=None)
 
-    @validator("ref")
+    @field_validator("ref")
     @skip_partial
     def check_ref(cls, ref, values):
-        validate_param_value(value=values.get("value"), ref=ref)
+        values = cls.get_data_from_values(values)
+        validate_param_value(value=cls.get_value_for_key("value", values), ref=ref)
         return ref
 
 

polyaxon/_flow/plugins/__init__.py

@@ -259,19 +259,25 @@ class V1Plugins(BaseSchemaModel):
 
     _IDENTIFIER = "plugins"
 
-    auth: Optional[BoolOrRef]
-    docker: Optional[BoolOrRef]
-    shm: Optional[BoolOrRef]
-    mount_artifacts_store: Optional[BoolOrRef] = Field(alias="mountArtifactsStore")
-    collect_artifacts: Optional[BoolOrRef] = Field(alias="collectArtifacts")
-    collect_logs: Optional[BoolOrRef] = Field(alias="collectLogs")
-    collect_resources: Optional[BoolOrRef] = Field(alias="collectResources")
-    sync_statuses: Optional[BoolOrRef] = Field(alias="syncStatuses")
-    auto_resume: Optional[BoolOrRef] = Field(alias="autoResume")
-    log_level: Optional[StrictStr] = Field(alias="logLevel")
-    external_host: Optional[BoolOrRef] = Field(alias="externalHost")
-    sidecar: Optional[Union[V1PolyaxonSidecarContainer, RefField]]
-    notifications: Optional[Union[List[V1Notification], RefField]]
+    auth: Optional[BoolOrRef] = None
+    docker: Optional[BoolOrRef] = None
+    shm: Optional[BoolOrRef] = None
+    mount_artifacts_store: Optional[BoolOrRef] = Field(
+        alias="mountArtifactsStore", default=None
+    )
+    collect_artifacts: Optional[BoolOrRef] = Field(
+        alias="collectArtifacts", default=None
+    )
+    collect_logs: Optional[BoolOrRef] = Field(alias="collectLogs", default=None)
+    collect_resources: Optional[BoolOrRef] = Field(
+        alias="collectResources", default=None
+    )
+    sync_statuses: Optional[BoolOrRef] = Field(alias="syncStatuses", default=None)
+    auto_resume: Optional[BoolOrRef] = Field(alias="autoResume", default=None)
+    log_level: Optional[StrictStr] = Field(alias="logLevel", default=None)
+    external_host: Optional[BoolOrRef] = Field(alias="externalHost", default=None)
+    sidecar: Optional[Union[V1PolyaxonSidecarContainer, RefField]] = None
+    notifications: Optional[Union[List[V1Notification], RefField]] = None
 
     @classmethod
     def get_or_create(

polyaxon/_flow/run/dag.py

@@ -7,7 +7,9 @@ from clipped.compact.pydantic import (
     PositiveInt,
     PrivateAttr,
     StrictStr,
-    validator,
+    field_validator,
+    validation_always,
+    validation_before,
 )
 from clipped.types.ref_or_obj import RefField
 
@@ -300,16 +302,17 @@ class V1Dag(BaseRun):
 
     _IDENTIFIER = V1RunKind.DAG
     _SWAGGER_FIELDS = ["volumes"]
+    _CUSTOM_DUMP_FIELDS = {"operations", "components", "environment"}
 
     kind: Literal[_IDENTIFIER] = _IDENTIFIER
-    operations: Optional[Union[List[V1Operation], RefField]]
-    components: Optional[Union[List[V1Component], RefField]]
-    environment: Optional[Union[V1Environment, RefField]]
-    connections: Optional[Union[List[StrictStr], RefField]]
-    volumes: Optional[Union[List[k8s_schemas.V1Volume], RefField]]
-    concurrency: Optional[Union[PositiveInt, RefField]]
+    operations: Optional[Union[List[V1Operation], RefField]] = None
+    components: Optional[Union[List[V1Component], RefField]] = None
+    environment: Optional[Union[V1Environment, RefField]] = None
+    connections: Optional[Union[List[StrictStr], RefField]] = None
+    volumes: Optional[Union[List[k8s_schemas.V1Volume], RefField]] = None
+    concurrency: Optional[Union[PositiveInt, RefField]] = None
     early_stopping: Optional[Union[List[V1EarlyStopping], RefField]] = Field(
-        alias="earlyStopping"
+        alias="earlyStopping", default=None
     )
 
     _dag: Dict[str, "DagOpSpec"] = PrivateAttr()
@@ -317,7 +320,7 @@ class V1Dag(BaseRun):
     _op_component_mapping: Dict[str, str] = PrivateAttr()
     _context: Dict[str, Any] = PrivateAttr()
 
-    @validator("volumes", always=True, pre=True)
+    @field_validator("volumes", **validation_always, **validation_before)
     def validate_volumes(cls, v):
         if not v:
             return v

polyaxon/_flow/run/dask/dask.py

@@ -124,12 +124,12 @@ class V1DaskJob(BaseRun, DestinationImageMixin):
     """
 
     _IDENTIFIER = V1RunKind.DASKJOB
-    _SWAGGER_FIELDS = ["volumes", "sidecars", "container"]
+    _CUSTOM_DUMP_FIELDS = {"job", "worker", "scheduler"}
 
     kind: Literal[_IDENTIFIER] = _IDENTIFIER
-    job: Optional[Union[V1DaskReplica, RefField]]
-    worker: Optional[Union[V1DaskReplica, RefField]]
-    scheduler: Optional[Union[V1DaskReplica, RefField]]
+    job: Optional[Union[V1DaskReplica, RefField]] = None
+    worker: Optional[Union[V1DaskReplica, RefField]] = None
+    scheduler: Optional[Union[V1DaskReplica, RefField]] = None
 
     def apply_image_destination(self, image: str):
         if self.job:

polyaxon/_flow/run/dask/replica.py

@@ -1,6 +1,11 @@
 from typing import List, Optional, Union
 
-from clipped.compact.pydantic import StrictStr, validator
+from clipped.compact.pydantic import (
+    StrictStr,
+    field_validator,
+    validation_always,
+    validation_before,
+)
 from clipped.types.ref_or_obj import IntOrRef, RefField
 
 from polyaxon._flow.environment import V1Environment
@@ -174,28 +179,29 @@ class V1DaskReplica(BaseSchemaModel):
 
     _IDENTIFIER = "replica"
     _SWAGGER_FIELDS = ["volumes", "sidecars", "container"]
+    _CUSTOM_DUMP_FIELDS = {"environment"}
 
-    replicas: Optional[IntOrRef]
-    environment: Optional[Union[V1Environment, RefField]]
-    connections: Optional[Union[List[StrictStr], RefField]]
-    volumes: Optional[Union[List[k8s_schemas.V1Volume], RefField]]
-    init: Optional[Union[List[V1Init], RefField]]
-    sidecars: Optional[Union[List[k8s_schemas.V1Container], RefField]]
-    container: Optional[Union[k8s_schemas.V1Container, RefField]]
+    replicas: Optional[IntOrRef] = None
+    environment: Optional[Union[V1Environment, RefField]] = None
+    connections: Optional[Union[List[StrictStr], RefField]] = None
+    volumes: Optional[Union[List[k8s_schemas.V1Volume], RefField]] = None
+    init: Optional[Union[List[V1Init], RefField]] = None
+    sidecars: Optional[Union[List[k8s_schemas.V1Container], RefField]] = None
+    container: Optional[Union[k8s_schemas.V1Container, RefField]] = None
 
-    @validator("volumes", always=True, pre=True)
+    @field_validator("volumes", **validation_always, **validation_before)
     def validate_volumes(cls, v):
         if not v:
             return v
         return [k8s_validation.validate_k8s_volume(vi) for vi in v]
 
-    @validator("sidecars", always=True, pre=True)
+    @field_validator("sidecars", **validation_always, **validation_before)
     def validate_helper_containers(cls, v):
         if not v:
             return v
         return [k8s_validation.validate_k8s_container(vi) for vi in v]
 
-    @validator("container", always=True, pre=True)
+    @field_validator("container", **validation_always, **validation_before)
     def validate_container(cls, v):
         return k8s_validation.validate_k8s_container(v)
 

polyaxon/_flow/run/job.py

@@ -1,7 +1,12 @@
 from typing import List, Optional, Union
 from typing_extensions import Literal
 
-from clipped.compact.pydantic import StrictStr, validator
+from clipped.compact.pydantic import (
+    StrictStr,
+    field_validator,
+    validation_always,
+    validation_before,
+)
 from clipped.types.ref_or_obj import RefField
 
 from polyaxon._flow.environment import V1Environment
@@ -213,28 +218,29 @@ class V1Job(BaseRun, DestinationImageMixin):
 
     _IDENTIFIER = V1RunKind.JOB
     _SWAGGER_FIELDS = ["volumes", "sidecars", "container"]
+    _CUSTOM_DUMP_FIELDS = {"environment", "init"}
 
     kind: Literal[_IDENTIFIER] = _IDENTIFIER
-    environment: Optional[Union[V1Environment, RefField]]
-    connections: Optional[Union[List[StrictStr], RefField]]
-    volumes: Optional[Union[List[k8s_schemas.V1Volume], RefField]]
-    init: Optional[Union[List[V1Init], RefField]]
-    sidecars: Optional[Union[List[k8s_schemas.V1Container], RefField]]
-    container: Optional[Union[k8s_schemas.V1Container, RefField]]
-
-    @validator("volumes", always=True, pre=True)
+    environment: Optional[Union[V1Environment, RefField]] = None
+    connections: Optional[Union[List[StrictStr], RefField]] = None
+    volumes: Optional[Union[List[k8s_schemas.V1Volume], RefField]] = None
+    init: Optional[Union[List[V1Init], RefField]] = None
+    sidecars: Optional[Union[List[k8s_schemas.V1Container], RefField]] = None
+    container: Optional[Union[k8s_schemas.V1Container, RefField]] = None
+
+    @field_validator("volumes", **validation_always, **validation_before)
     def validate_volumes(cls, v):
         if not v:
             return v
         return [k8s_validation.validate_k8s_volume(vi) for vi in v]
 
-    @validator("sidecars", always=True, pre=True)
+    @field_validator("sidecars", **validation_always, **validation_before)
     def validate_helper_containers(cls, v):
         if not v:
             return v
         return [k8s_validation.validate_k8s_container(vi) for vi in v]
 
-    @validator("container", always=True, pre=True)
+    @field_validator("container", **validation_always, **validation_before)
     def validate_container(cls, v):
         return k8s_validation.validate_k8s_container(v)
 

polyaxon/_flow/run/kubeflow/mpi_job.py

@@ -133,13 +133,18 @@ class V1MPIJob(BaseRun, DestinationImageMixin):
     """
 
     _IDENTIFIER = V1RunKind.MPIJOB
+    _CUSTOM_DUMP_FIELDS = {"launcher", "worker"}
 
     kind: Literal[_IDENTIFIER] = _IDENTIFIER
-    clean_pod_policy: Optional[V1CleanPodPolicy] = Field(alias="cleanPodPolicy")
-    scheduling_policy: Optional[V1SchedulingPolicy] = Field(alias="schedulingPolicy")
-    slots_per_worker: Optional[IntOrRef] = Field(alias="slotsPerWorker")
-    launcher: Optional[Union[V1KFReplica, RefField]]
-    worker: Optional[Union[V1KFReplica, RefField]]
+    clean_pod_policy: Optional[V1CleanPodPolicy] = Field(
+        alias="cleanPodPolicy", default=None
+    )
+    scheduling_policy: Optional[V1SchedulingPolicy] = Field(
+        alias="schedulingPolicy", default=None
+    )
+    slots_per_worker: Optional[IntOrRef] = Field(alias="slotsPerWorker", default=None)
+    launcher: Optional[Union[V1KFReplica, RefField]] = None
+    worker: Optional[Union[V1KFReplica, RefField]] = None
 
     def apply_image_destination(self, image: str):
         if self.launcher:

polyaxon/_flow/run/kubeflow/mx_job.py

@@ -202,17 +202,33 @@ class V1MXJob(BaseRun, DestinationImageMixin):
     """
 
     _IDENTIFIER = V1RunKind.MXJOB
+    _CUSTOM_DUMP_FIELDS = {
+        "scheduler",
+        "server",
+        "worker",
+        "tuner",
+        "tunerTracker",
+        "tunerServer",
+    }
 
     kind: Literal[_IDENTIFIER] = _IDENTIFIER
-    clean_pod_policy: Optional[V1CleanPodPolicy] = Field(alias="cleanPodPolicy")
-    scheduling_policy: Optional[V1SchedulingPolicy] = Field(alias="schedulingPolicy")
-    mode: Optional[MXJobMode]
-    scheduler: Optional[Union[V1KFReplica, RefField]]
-    server: Optional[Union[V1KFReplica, RefField]]
-    worker: Optional[Union[V1KFReplica, RefField]]
-    tuner: Optional[Union[V1KFReplica, RefField]]
-    tuner_tracker: Optional[Union[V1KFReplica, RefField]] = Field(alias="tunerTracker")
-    tuner_server: Optional[Union[V1KFReplica, RefField]] = Field(alias="tunerServer")
+    clean_pod_policy: Optional[V1CleanPodPolicy] = Field(
+        alias="cleanPodPolicy", default=None
+    )
+    scheduling_policy: Optional[V1SchedulingPolicy] = Field(
+        alias="schedulingPolicy", default=None
+    )
+    mode: Optional[MXJobMode] = None
+    scheduler: Optional[Union[V1KFReplica, RefField]] = None
+    server: Optional[Union[V1KFReplica, RefField]] = None
+    worker: Optional[Union[V1KFReplica, RefField]] = None
+    tuner: Optional[Union[V1KFReplica, RefField]] = None
+    tuner_tracker: Optional[Union[V1KFReplica, RefField]] = Field(
+        alias="tunerTracker", default=None
+    )
+    tuner_server: Optional[Union[V1KFReplica, RefField]] = Field(
+        alias="tunerServer", default=None
+    )
 
     def apply_image_destination(self, image: str):
         if self.scheduler:

polyaxon/_flow/run/kubeflow/paddle_job.py

@@ -27,10 +27,10 @@ class V1PaddleElasticPolicy(BaseSchemaModel):
 
     _IDENTIFIER = "elasticPolicy"
 
-    min_replicas: Optional[IntOrRef] = Field(alias="minReplicas")
-    max_replicas: Optional[IntOrRef] = Field(alias="maxReplicas")
-    max_restarts: Optional[IntOrRef] = Field(alias="maxRestarts")
-    metrics: Optional[List[Dict]] = Field(alias="Metrics")
+    min_replicas: Optional[IntOrRef] = Field(alias="minReplicas", default=None)
+    max_replicas: Optional[IntOrRef] = Field(alias="maxReplicas", default=None)
+    max_restarts: Optional[IntOrRef] = Field(alias="maxRestarts", default=None)
+    metrics: Optional[List[Dict]] = Field(alias="Metrics", default=None)
 
 
 class V1PaddleJob(BaseRun, DestinationImageMixin):
@@ -151,13 +151,20 @@ class V1PaddleJob(BaseRun, DestinationImageMixin):
     """
 
     _IDENTIFIER = V1RunKind.PADDLEJOB
+    _CUSTOM_DUMP_FIELDS = {"master", "worker"}
 
     kind: Literal[_IDENTIFIER] = _IDENTIFIER
-    clean_pod_policy: Optional[V1CleanPodPolicy] = Field(alias="cleanPodPolicy")
-    scheduling_policy: Optional[V1SchedulingPolicy] = Field(alias="schedulingPolicy")
-    elastic_policy: Optional[V1PaddleElasticPolicy] = Field(alias="elasticPolicy")
-    master: Optional[Union[V1KFReplica, RefField]]
-    worker: Optional[Union[V1KFReplica, RefField]]
+    clean_pod_policy: Optional[V1CleanPodPolicy] = Field(
+        alias="cleanPodPolicy", default=None
+    )
+    scheduling_policy: Optional[V1SchedulingPolicy] = Field(
+        alias="schedulingPolicy", default=None
+    )
+    elastic_policy: Optional[V1PaddleElasticPolicy] = Field(
+        alias="elasticPolicy", default=None
+    )
+    master: Optional[Union[V1KFReplica, RefField]] = None
+    worker: Optional[Union[V1KFReplica, RefField]] = None
 
     def apply_image_destination(self, image: str):
         if self.master:

polyaxon/_flow/run/kubeflow/pytorch_job.py

@@ -34,17 +34,17 @@ class V1PytorchElasticPolicy(BaseSchemaModel):
 
     _IDENTIFIER = "elasticPolicy"
 
-    min_replicas: Optional[IntOrRef] = Field(alias="minReplicas")
-    max_replicas: Optional[IntOrRef] = Field(alias="maxReplicas")
-    rdvz_backend: Optional[StrictStr] = Field(alias="rdvzBackend")
-    rdvz_port: Optional[IntOrRef] = Field(alias="rdvzPort")
-    rdvz_host: Optional[StrictStr] = Field(alias="rdvzHost")
-    rdvz_id: Optional[StrictStr] = Field(alias="rdvzId")
-    rdvz_conf: Optional[List[Dict]] = Field(alias="rdvzConf")
-    standalone: Optional[BoolOrRef]
-    n_proc_per_node: Optional[IntOrRef] = Field(alias="nProcPerNode")
-    max_restarts: Optional[IntOrRef] = Field(alias="maxRestarts")
-    metrics: Optional[List[Dict]] = Field(alias="Metrics")
+    min_replicas: Optional[IntOrRef] = Field(alias="minReplicas", default=None)
+    max_replicas: Optional[IntOrRef] = Field(alias="maxReplicas", default=None)
+    rdvz_backend: Optional[StrictStr] = Field(alias="rdvzBackend", default=None)
+    rdvz_port: Optional[IntOrRef] = Field(alias="rdvzPort", default=None)
+    rdvz_host: Optional[StrictStr] = Field(alias="rdvzHost", default=None)
+    rdvz_id: Optional[StrictStr] = Field(alias="rdvzId", default=None)
+    rdvz_conf: Optional[List[Dict]] = Field(alias="rdvzConf", default=None)
+    standalone: Optional[BoolOrRef] = None
+    n_proc_per_node: Optional[IntOrRef] = Field(alias="nProcPerNode", default=None)
+    max_restarts: Optional[IntOrRef] = Field(alias="maxRestarts", default=None)
+    metrics: Optional[List[Dict]] = Field(alias="Metrics", default=None)
 
 
 class V1PytorchJob(BaseRun, DestinationImageMixin):
@@ -163,14 +163,21 @@ class V1PytorchJob(BaseRun, DestinationImageMixin):
     """
 
     _IDENTIFIER = V1RunKind.PYTORCHJOB
+    _CUSTOM_DUMP_FIELDS = {"master", "worker"}
 
     kind: Literal[_IDENTIFIER] = _IDENTIFIER
-    clean_pod_policy: Optional[V1CleanPodPolicy] = Field(alias="cleanPodPolicy")
-    scheduling_policy: Optional[V1SchedulingPolicy] = Field(alias="schedulingPolicy")
-    elastic_policy: Optional[V1PytorchElasticPolicy] = Field(alias="elasticPolicy")
-    n_proc_per_node: Optional[IntOrRef] = Field(alias="nProcPerNode")
-    master: Optional[Union[V1KFReplica, RefField]]
-    worker: Optional[Union[V1KFReplica, RefField]]
+    clean_pod_policy: Optional[V1CleanPodPolicy] = Field(
+        alias="cleanPodPolicy", default=None
+    )
+    scheduling_policy: Optional[V1SchedulingPolicy] = Field(
+        alias="schedulingPolicy", default=None
+    )
+    elastic_policy: Optional[V1PytorchElasticPolicy] = Field(
+        alias="elasticPolicy", default=None
+    )
+    n_proc_per_node: Optional[IntOrRef] = Field(alias="nProcPerNode", default=None)
+    master: Optional[Union[V1KFReplica, RefField]] = None
+    worker: Optional[Union[V1KFReplica, RefField]] = None
 
     def apply_image_destination(self, image: str):
         if self.master:

polyaxon/_flow/run/kubeflow/replica.py

@@ -1,6 +1,11 @@
 from typing import List, Optional, Union
 
-from clipped.compact.pydantic import StrictStr, validator
+from clipped.compact.pydantic import (
+    StrictStr,
+    field_validator,
+    validation_always,
+    validation_before,
+)
 from clipped.types.ref_or_obj import IntOrRef, RefField
 
 from polyaxon._flow.environment import V1Environment
@@ -202,28 +207,29 @@ class V1KFReplica(BaseSchemaModel):
 
     _IDENTIFIER = "replica"
     _SWAGGER_FIELDS = ["volumes", "sidecars", "container"]
+    _CUSTOM_DUMP_FIELDS = {"environment"}
 
-    replicas: Optional[IntOrRef]
-    environment: Optional[Union[V1Environment, RefField]]
-    connections: Optional[Union[List[StrictStr], RefField]]
-    volumes: Optional[Union[List[k8s_schemas.V1Volume], RefField]]
-    init: Optional[Union[List[V1Init], RefField]]
-    sidecars: Optional[Union[List[k8s_schemas.V1Container], RefField]]
-    container: Optional[Union[k8s_schemas.V1Container, RefField]]
+    replicas: Optional[IntOrRef] = None
+    environment: Optional[Union[V1Environment, RefField]] = None
+    connections: Optional[Union[List[StrictStr], RefField]] = None
+    volumes: Optional[Union[List[k8s_schemas.V1Volume], RefField]] = None
+    init: Optional[Union[List[V1Init], RefField]] = None
+    sidecars: Optional[Union[List[k8s_schemas.V1Container], RefField]] = None
+    container: Optional[Union[k8s_schemas.V1Container, RefField]] = None
 
-    @validator("volumes", always=True, pre=True)
+    @field_validator("volumes", **validation_always, **validation_before)
     def validate_volumes(cls, v):
         if not v:
             return v
         return [k8s_validation.validate_k8s_volume(vi) for vi in v]
 
-    @validator("sidecars", always=True, pre=True)
+    @field_validator("sidecars", **validation_always, **validation_before)
     def validate_helper_containers(cls, v):
         if not v:
             return v
         return [k8s_validation.validate_k8s_container(vi) for vi in v]
 
-    @validator("container", always=True, pre=True)
+    @field_validator("container", **validation_always, **validation_before)
     def validate_container(cls, v):
         return k8s_validation.validate_k8s_container(v)