polyaxon 2.6.0rc1__tar.gz → 2.7.0rc0__tar.gz

{polyaxon-2.6.0rc1 → polyaxon-2.7.0rc0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: polyaxon
-Version: 2.6.0rc1
+Version: 2.7.0rc0
 Summary: Command Line Interface (CLI) and client to interact with Polyaxon API.
 Home-page: https://github.com/polyaxon/polyaxon
 Author: Polyaxon, Inc.
@@ -51,7 +51,7 @@ Provides-Extra: sandbox
 [![Slack](https://img.shields.io/badge/Slack-1.5k%20members-blue.svg?style=flat&logo=slack&longCache=true)](https://polyaxon.com/slack/)
 
 [![Docs](https://img.shields.io/badge/docs-stable-brightgreen.svg?style=flat&longCache=true)](https://polyaxon.com/docs/)
-[![Release](https://img.shields.io/badge/release-v2.6.0-brightgreen.svg?longCache=true)](https://polyaxon.com/docs/releases/2-6/)
+[![Release](https://img.shields.io/badge/release-v2.7.0-brightgreen.svg?longCache=true)](https://polyaxon.com/docs/releases/2-7/)
 [![GitHub](https://img.shields.io/badge/issue_tracker-github-blue?style=flat&logo=github&longCache=true)](https://github.com/polyaxon/polyaxon/issues)
 [![GitHub](https://img.shields.io/badge/roadmap-github-blue?style=flat&logo=github&longCache=true)](https://github.com/orgs/polyaxon/projects/5)
 

{polyaxon-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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-2.6.0rc1 → polyaxon-2.7.0rc0}/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