polyaxon 2.6.0rc1__py3-none-any.whl → 2.7.0rc0__py3-none-any.whl

polyaxon/_cli/operations.py

@@ -1466,7 +1466,7 @@ def logs(ctx, project, uid, follow, hide_time, all_containers, all_info, offline
             )
 
         try:
-            from traceml.logging.streamer import get_logs_streamer, load_logs_from_path
+            from traceml.logging.streamer import load_logs_from_path
 
             load_logs_from_path(
                 logs_path=logs_path,

polyaxon/_cli/projects.py

@@ -46,7 +46,7 @@ def project(ctx, _project):  # pylint:disable=redefined-outer-name
 @click.option(
     *OPTIONS_NAME["args"],
     type=str,
-    help="The project name, e.g. 'mnist' or 'acme/mnist'."
+    help="The project name, e.g. 'mnist' or 'acme/mnist'.",
 )
 @click.option("--description", type=str, help="Description of the project.")
 @click.option("--tags", type=str, help="Tags of the project (comma separated values).")
@@ -287,7 +287,7 @@ def delete(ctx, _project, yes):
 @click.option(
     *OPTIONS_NAME["args"],
     type=str,
-    help="Name of the project, must be unique for the same user."
+    help="Name of the project, must be unique for the same user.",
 )
 @click.option("--description", type=str, help="Description of the project.")
 @click.option("--tags", type=str, help="Tags of the project (comma separated values).")

polyaxon/_compiler/contexts/contexts.py

@@ -98,24 +98,24 @@ def resolve_globals_contexts(
         run_outputs_path = ctx_paths.CONTEXT_MOUNT_RUN_OUTPUTS_FORMAT.format(run_path)
         resolved_contexts[ctx_sections.GLOBALS][
             ctx_keys.RUN_ARTIFACTS_PATH
-        ] = run_artifacts_path
+        ] = run_artifacts_path  # fmt: skip
         resolved_contexts[ctx_sections.GLOBALS][
             ctx_keys.RUN_OUTPUTS_PATH
-        ] = run_outputs_path
+        ] = run_outputs_path  # fmt: skip
     elif artifacts_store:
         run_artifacts_path = os.path.join(artifacts_store.store_path, run_path)
         run_outputs_path = os.path.join(run_artifacts_path, "outputs")
         resolved_contexts[ctx_sections.GLOBALS][
             ctx_keys.RUN_ARTIFACTS_PATH
-        ] = run_artifacts_path
+        ] = run_artifacts_path  # fmt: skip
         resolved_contexts[ctx_sections.GLOBALS][
             ctx_keys.RUN_OUTPUTS_PATH
-        ] = run_outputs_path
+        ] = run_outputs_path  # fmt: skip
 
     if plugins.mount_artifacts_store and artifacts_store:
         resolved_contexts[ctx_sections.GLOBALS][
             ctx_keys.STORE_PATH
-        ] = artifacts_store.store_path
+        ] = artifacts_store.store_path  # fmt: skip
     return resolved_contexts
 
 

polyaxon/_compiler/contexts/dask_job.py

@@ -21,7 +21,7 @@ class DaskJobContextsManager(BaseContextsManager):
     ) -> Dict:
         contexts["init"] = {}
         contexts["connections"] = {}
-        job = compiled_operation.run  # type: V1DaskJob
+        job: V1DaskJob = compiled_operation.run
 
         def _get_replica(replica: Optional[V1DaskReplica]) -> Dict:
             if not replica:

polyaxon/_compiler/contexts/kubeflow/mpi_job.py

@@ -21,7 +21,7 @@ class MPIJobContextsManager(BaseContextsManager):
     ) -> Dict:
         contexts["init"] = {}
         contexts["connections"] = {}
-        job = compiled_operation.run  # type: V1MPIJob
+        job: V1MPIJob = compiled_operation.run
 
         def _get_replica(replica: Optional[V1KFReplica]) -> Dict:
             if not replica:

polyaxon/_compiler/contexts/kubeflow/mx_job.py

@@ -21,7 +21,7 @@ class MXJobContextsManager(BaseContextsManager):
     ) -> Dict:
         contexts["init"] = {}
         contexts["connections"] = {}
-        job = compiled_operation.run  # type: V1MXJob
+        job: V1MXJob = compiled_operation.run
 
         def _get_replica(replica: Optional[V1KFReplica]) -> Dict:
             if not replica:

polyaxon/_compiler/contexts/kubeflow/paddle_job.py

@@ -21,7 +21,7 @@ class PaddleJobContextsManager(BaseContextsManager):
     ) -> Dict:
         contexts["init"] = {}
         contexts["connections"] = {}
-        job = compiled_operation.run  # type: V1PaddleJob
+        job: V1PaddleJob = compiled_operation.run
 
         def _get_replica(replica: Optional[V1KFReplica]) -> Dict:
             if not replica:

polyaxon/_compiler/contexts/kubeflow/pytroch_job.py

@@ -21,7 +21,7 @@ class PytorchJobContextsManager(BaseContextsManager):
     ) -> Dict:
         contexts["init"] = {}
         contexts["connections"] = {}
-        job = compiled_operation.run  # type: V1PytorchJob
+        job: V1PytorchJob = compiled_operation.run
 
         def _get_replica(replica: Optional[V1KFReplica]) -> Dict:
             if not replica:

polyaxon/_compiler/contexts/kubeflow/tf_job.py

@@ -21,7 +21,7 @@ class TfJobContextsManager(BaseContextsManager):
     ) -> Dict:
         contexts["init"] = {}
         contexts["connections"] = {}
-        job = compiled_operation.run  # type: V1TFJob
+        job: V1TFJob = compiled_operation.run
 
         def _get_replica(replica: Optional[V1KFReplica]) -> Dict:
             if not replica:

polyaxon/_compiler/contexts/kubeflow/xgb_job.py

@@ -21,7 +21,7 @@ class XGBoostJobContextsManager(BaseContextsManager):
     ) -> Dict:
         contexts["init"] = {}
         contexts["connections"] = {}
-        job = compiled_operation.run  # type: V1XGBoostJob
+        job: V1XGBoostJob = compiled_operation.run
 
         def _get_replica(replica: Optional[V1KFReplica]) -> Dict:
             if not replica:

polyaxon/_compiler/contexts/ray_job.py

@@ -21,7 +21,7 @@ class RayJobContextsManager(BaseContextsManager):
     ) -> Dict:
         contexts["init"] = {}
         contexts["connections"] = {}
-        job = compiled_operation.run  # type: V1RayJob
+        job: V1RayJob = compiled_operation.run
 
         def _get_replica(replica: Optional[V1RayReplica]) -> Dict:
             if not replica:

polyaxon/_compiler/resolver/agent.py

@@ -71,7 +71,7 @@ class AgentResolver(BaseSchemaModel):
         if agent_config.artifacts_store:  # Resolve default artifacts store
             self.connection_by_names[
                 agent_config.artifacts_store.name
-            ] = agent_config.artifacts_store
+            ] = agent_config.artifacts_store  # fmt: skip
 
         if (
             compiled_operation.is_job_run

polyaxon/_compiler/resolver/runtime.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import Any, Dict, List, Optional
+from typing import Dict, List, Optional
 
 from polyaxon import settings
 from polyaxon._compiler.contexts import resolve_contexts, resolve_globals_contexts

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/_deploy/schemas/ui.py

@@ -13,4 +13,6 @@ class UIConfig(BaseSchemaModel):
     assets_version: Optional[StrictStr] = Field(alias="assetsVersion", default=None)
     admin_enabled: Optional[bool] = Field(alias="adminEnabled", default=None)
     single_url: Optional[bool] = Field(alias="singleUrl", default=None)
-    default_streams_url: Optional[StrictStr] = Field(alias="defaultStreamsUrl", default=None)
+    default_streams_url: Optional[StrictStr] = Field(
+        alias="defaultStreamsUrl", default=None
+    )

polyaxon/_docker/converter/base/base.py

@@ -162,7 +162,7 @@ class BaseConverter(
         if memory:
             docker_resources[
                 "memory"
-            ] = docker_types.V1ResourceRequirements.from_k8s_memory(memory)
+            ] = docker_types.V1ResourceRequirements.from_k8s_memory(memory)  # fmt: skip
         if gpus:
             docker_resources["gpus"] = gpus
         return docker_types.V1ResourceRequirements.from_dict(docker_resources)

polyaxon/_docker/converter/base/env_vars.py

@@ -98,7 +98,7 @@ class EnvMixin(BaseConverter):
 
         try:
             secret_value = orjson_loads(resource.name)
-        except orjson.JSONDecodeError as e:
+        except orjson.JSONDecodeError:
             return items_from
 
         for item in resource.items:

polyaxon/_docker/converter/converters/job.py

@@ -18,7 +18,7 @@ class JobConverter(JobMixin, BaseConverter):
         default_sa: Optional[str] = None,
         default_auth: bool = False,
     ) -> List[docker_types.V1Container]:
-        job = compiled_operation.run  # type: V1Job
+        job: V1Job = compiled_operation.run
         plugins = V1Plugins.get_or_create(
             config=compiled_operation.plugins, auth=default_auth
         )

polyaxon/_docker/converter/converters/service.py

@@ -4,7 +4,7 @@ from polyaxon._connections import V1Connection, V1ConnectionResource
 from polyaxon._docker import docker_types
 from polyaxon._docker.converter.base import BaseConverter
 from polyaxon._docker.converter.mixins import ServiceMixin
-from polyaxon._flow import V1CompiledOperation, V1Plugins, V1Service
+from polyaxon._flow import V1CompiledOperation, V1Plugins
 
 
 class ServiceConverter(ServiceMixin, BaseConverter):

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/environment/__init__.py

@@ -8,6 +8,7 @@ from clipped.compact.pydantic import (
     validation_always,
     validation_before,
 )
+from clipped.types import GenericStr
 
 from polyaxon._k8s import k8s_schemas, k8s_validation
 from polyaxon._schemas.base import BaseSchemaModel
@@ -395,9 +396,9 @@ class V1Environment(BaseSchemaModel):
         "dnsConfig",
     ]
 
-    labels: Optional[Dict[StrictStr, StrictStr]] = None
-    annotations: Optional[Dict[StrictStr, StrictStr]] = None
-    node_selector: Optional[Dict[StrictStr, StrictStr]] = Field(
+    labels: Optional[Dict[StrictStr, GenericStr]] = None
+    annotations: Optional[Dict[StrictStr, GenericStr]] = None
+    node_selector: Optional[Dict[StrictStr, GenericStr]] = Field(
         alias="nodeSelector", default=None
     )
     affinity: Optional[Union[k8s_schemas.V1Affinity, Dict]] = None
@@ -428,7 +429,7 @@ class V1Environment(BaseSchemaModel):
     priority: Optional[int] = None
     restart_policy: Optional[
         Literal["Always", "OnFailure", "Never", "ExitCode"]
-    ] = Field(alias="restartPolicy", default=None)
+    ] = Field(alias="restartPolicy", default=None)  # fmt: skip
 
     @field_validator("affinity", **validation_always, **validation_before)
     def validate_affinity(cls, v):

polyaxon/_flow/io/io.py

@@ -46,6 +46,8 @@ def validate_io_value(
             default=default,
         )
         if validation:
+            if isinstance(validation, dict):
+                validation = V1Validation(**validation)
             validation.run_validation(
                 value=parsed_value, type=type, is_optional=is_optional
             )

polyaxon/_flow/matrix/params.py

@@ -1039,7 +1039,7 @@ class V1HpQLogUniform(BaseHpParamConfig):
     kind: Literal[_IDENTIFIER] = _IDENTIFIER
     value: Optional[
         Union[QLogUniform, QLogUniformList, QLogUniformStr, RefField]
-    ] = None
+    ] = None  # fmt: skip
 
     @property
     def is_distribution(self):

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

@@ -8,7 +8,6 @@ from clipped.compact.pydantic import (
     field_validator,
     model_validator,
     validation_after,
-    validation_before,
 )
 from clipped.config.patch_strategy import PatchStrategy
 from clipped.config.schema import skip_partial, to_partial
@@ -649,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/params/params.py

@@ -386,7 +386,7 @@ class V1Param(BaseSchemaModel, ctx_refs.RefMixin, ParamValueMixin):
 
     _IDENTIFIER = "param"
 
-    value: Any
+    value: Optional[Any] = None
     ref: Optional[StrictStr] = None
     context_only: Optional[bool] = Field(alias="contextOnly", default=False)
     connection: Optional[StrictStr] = None
@@ -479,7 +479,7 @@ class ParamSpec(
         raise PolyaxonValidationError(
             "Param `{}` with type `{}`, "
             "cannot be turned to an init container automatically.".format(
-                self.name, self.type, self.param.ref
+                self.name, self.type
             )
         )
 

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/dag.py

@@ -444,7 +444,7 @@ class V1Dag(BaseRun):
         for g_context in ctx_sections.GLOBALS_CONTEXTS:
             self._context[
                 "dag.{}.{}".format(ctx_sections.GLOBALS, g_context)
-            ] = V1IO.construct(name=g_context, type="str", value="", is_optional=True)
+            ] = V1IO.construct(name=g_context, type="str", value="", is_optional=True)  # fmt: skip
 
         self._context["dag.{}".format(ctx_sections.INPUTS)] = V1IO.construct(
             name="inputs", type="dict", value={}, is_optional=True
@@ -541,19 +541,19 @@ class V1Dag(BaseRun):
             # We allow to resolve name, status, project, all outputs/inputs, iteration
             self._context[
                 "ops.{}.{}".format(op_name, ctx_sections.INPUTS)
-            ] = V1IO.construct(name="inputs", type="dict", value={}, is_optional=True)
+            ] = V1IO.construct(name="inputs", type="dict", value={}, is_optional=True)  # fmt: skip
             self._context[
                 "ops.{}.{}".format(op_name, ctx_sections.OUTPUTS)
-            ] = V1IO.construct(name="outputs", type="dict", value={}, is_optional=True)
+            ] = V1IO.construct(name="outputs", type="dict", value={}, is_optional=True)  # fmt: skip
             self._context[
                 "ops.{}.{}".format(op_name, ctx_sections.GLOBALS)
-            ] = V1IO.construct(name="globals", type="str", value="", is_optional=True)
+            ] = V1IO.construct(name="globals", type="str", value="", is_optional=True)  # fmt: skip
             self._context[
                 "ops.{}.{}".format(op_name, ctx_sections.ARTIFACTS)
-            ] = V1IO.construct(name="artifacts", type="str", value="", is_optional=True)
+            ] = V1IO.construct(name="artifacts", type="str", value="", is_optional=True)  # fmt: skip
             self._context[
                 "ops.{}.{}".format(op_name, ctx_sections.INPUTS_OUTPUTS)
-            ] = V1IO.construct(name="io", type="str", value={}, is_optional=True)
+            ] = V1IO.construct(name="io", type="str", value={}, is_optional=True)  # fmt: skip
 
         for op in self.operations:
             if op.has_component_reference:

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: