polyaxon 2.0.0rc49__py3-none-any.whl → 2.4.0rc1__py3-none-any.whl

polyaxon/_local_process/agent.py

@@ -0,0 +1,6 @@
+from polyaxon._local_process.executor import Executor
+from polyaxon._runner.agent.sync_agent import BaseSyncAgent
+
+
+class Agent(BaseSyncAgent):
+    EXECUTOR = Executor

polyaxon/_local_process/converter/__init__.py

@@ -0,0 +1 @@
+

polyaxon/_local_process/converter/base/__init__.py

@@ -0,0 +1 @@
+from polyaxon._local_process.converter.base.base import BaseConverter

polyaxon/_local_process/converter/base/base.py

@@ -0,0 +1,140 @@
+from typing import Dict, Iterable, List, Optional, Union
+
+from polyaxon import settings
+from polyaxon._auxiliaries import V1PolyaxonSidecarContainer
+from polyaxon._connections import V1Connection, V1ConnectionResource
+from polyaxon._flow import V1Environment, V1Init, V1Plugins
+from polyaxon._k8s import k8s_schemas
+from polyaxon._local_process import process_types
+from polyaxon._local_process.converter.base.containers import ContainerMixin
+from polyaxon._local_process.converter.base.env_vars import EnvMixin
+from polyaxon._local_process.converter.base.init import InitConverter
+from polyaxon._local_process.converter.base.main import MainConverter
+from polyaxon._local_process.converter.base.mounts import MountsMixin
+from polyaxon._runner.converter import BaseConverter as _BaseConverter
+from polyaxon._runner.kinds import RunnerKind
+from polyaxon.exceptions import PolyaxonConverterError
+
+
+class BaseConverter(
+    MainConverter, InitConverter, ContainerMixin, EnvMixin, MountsMixin, _BaseConverter
+):
+    RUNNER_KIND = RunnerKind.PROCESS
+
+    @classmethod
+    def _get_sidecar_container(
+        cls,
+        container_id: str,
+        polyaxon_sidecar: V1PolyaxonSidecarContainer,
+        env: List[process_types.V1EnvVar],
+        artifacts_store: V1Connection,
+        plugins: V1Plugins,
+        run_path: Optional[str],
+    ) -> Optional[process_types.V1Container]:
+        return None
+
+    @classmethod
+    def _k8s_to_process_env_var(
+        cls,
+        env_var: List[k8s_schemas.V1EnvVar],
+    ) -> List[process_types.V1EnvVar]:
+        if not env_var:
+            return []
+
+        process_env_var = []
+        for item in env_var:
+            if isinstance(item, dict):
+                try:
+                    item = k8s_schemas.V1EnvVar(**item)
+                except (ValueError, TypeError) as e:
+                    raise PolyaxonConverterError(
+                        f"Could not parse env var value `{item}` from the K8S schema in container section"
+                    ) from e
+            process_env_var.append(cls._get_env_var(name=item.name, value=item.value))
+
+        return process_env_var
+
+    @staticmethod
+    def _new_container(name: str) -> process_types.V1Container:
+        return process_types.V1Container(name=name)
+
+    @classmethod
+    def _ensure_container(
+        cls,
+        container: Union[k8s_schemas.V1Container, process_types.V1Container],
+        volumes: List[k8s_schemas.V1Volume],
+    ) -> process_types.V1Container:
+        if not isinstance(container, k8s_schemas.V1Container):
+            return container
+        return process_types.V1Container(
+            name=container.name,
+            command=container.command,
+            args=container.args,
+            env=cls._k8s_to_process_env_var(container.env),
+            working_dir=container.working_dir,
+        )
+
+    def get_replica_resource(
+        self,
+        environment: V1Environment,
+        plugins: V1Plugins,
+        volumes: List[k8s_schemas.V1Volume],
+        init: List[V1Init],
+        sidecars: List[k8s_schemas.V1Container],
+        container: k8s_schemas.V1Container,
+        artifacts_store: V1Connection,
+        connections: List[str],
+        connection_by_names: Dict[str, V1Connection],
+        secrets: Optional[Iterable[V1ConnectionResource]],
+        config_maps: Optional[Iterable[V1ConnectionResource]],
+        kv_env_vars: List[List],
+        default_sa: Optional[str] = None,
+        ports: List[int] = None,
+    ) -> List[process_types.V1Container]:
+        volumes = volumes or []
+        init = init or []
+        sidecars = sidecars or []
+        connections = connections or []
+        environment = environment or V1Environment()
+        environment.service_account_name = (
+            environment.service_account_name
+            or default_sa
+            or settings.AGENT_CONFIG.runs_sa
+        )
+
+        init_connections = self.filter_connections_from_init(init=init)
+
+        init_containers = self.get_init_containers(
+            polyaxon_init=self.polyaxon_init,
+            plugins=plugins,
+            artifacts_store=artifacts_store,
+            init_connections=init_connections,
+            init_containers=self.filter_containers_from_init(init=init),
+            connection_by_names=connection_by_names,
+            log_level=plugins.log_level,
+            volumes=volumes,
+        )
+
+        sidecar_containers = self.get_sidecar_containers(
+            polyaxon_sidecar=self.polyaxon_sidecar,
+            plugins=plugins,
+            artifacts_store=artifacts_store,
+            sidecar_containers=sidecars,
+            log_level=plugins.log_level,
+            volumes=volumes,
+        )
+
+        main_container = self.get_main_container(
+            main_container=self._ensure_container(container, volumes=volumes),
+            plugins=plugins,
+            artifacts_store=artifacts_store,
+            connections=connections,
+            init_connections=init_connections,
+            connection_by_names=connection_by_names,
+            secrets=secrets,
+            config_maps=config_maps,
+            ports=ports,
+            kv_env_vars=kv_env_vars,
+        )
+
+        return init_containers + sidecar_containers + [main_container]

polyaxon/_local_process/converter/base/containers.py

@@ -0,0 +1,69 @@
+from typing import Any, Dict, List, Optional
+
+from clipped.utils.lists import to_list
+from clipped.utils.sanitizers import sanitize_value
+
+from polyaxon._containers.names import sanitize_container_name
+from polyaxon._local_process import process_types
+from polyaxon._runner.converter import BaseConverter
+from polyaxon._runner.converter.common.containers import sanitize_container_command_args
+
+
+class ContainerMixin(BaseConverter):
+    @classmethod
+    def _patch_container(
+        cls,
+        container: process_types.V1Container,
+        name: Optional[str] = None,
+        command: Optional[List[str]] = None,
+        args: Optional[List[str]] = None,
+        env: Optional[List[process_types.V1EnvVar]] = None,
+        env_from: Optional[List[Any]] = None,
+        **kwargs,
+    ) -> process_types.V1Container:
+        container.name = sanitize_container_name(name or container.name)
+        container.env = to_list(container.env, check_none=True) + to_list(
+            env, check_none=True
+        )
+        if not any([container.command, container.args]):
+            container.command = command
+            container.args = args
+
+        return cls._sanitize_container(container)
+
+    @staticmethod
+    def _sanitize_container_env(
+        env: List[process_types.V1EnvVar],
+    ) -> List[process_types.V1EnvVar]:
+        def sanitize_env_dict(d: Dict):
+            return process_types.V1EnvVar(
+                __root__={
+                    d_k: sanitize_value(d_v, handle_dict=False)
+                    for d_k, d_v in d.items()
+                }
+            )
+
+        results = []
+        for e in env or []:
+            if isinstance(e, dict):
+                e = sanitize_env_dict(e)
+                results.append(e)
+            elif isinstance(e, tuple):
+                if e[1] is not None:
+                    e = process_types.V1EnvVar(
+                        __root__=(e[0], sanitize_value(e[1], handle_dict=False))
+                    )
+                results.append(e)
+            elif isinstance(e, process_types.V1EnvVar):
+                results.append(e)
+
+        return results
+
+    @classmethod
+    def _sanitize_container(
+        cls,
+        container: process_types.V1Container,
+    ) -> process_types.V1Container:
+        container = sanitize_container_command_args(container)
+        container.env = cls._sanitize_container_env(container.env)
+        return container

polyaxon/_local_process/converter/base/env_vars.py

@@ -0,0 +1,253 @@
+import os
+
+from typing import Any, Iterable, List, Optional
+
+import orjson
+
+from clipped.utils.enums import get_enum_value
+from clipped.utils.json import orjson_dumps, orjson_loads
+from clipped.utils.lists import to_list
+
+from polyaxon import settings
+from polyaxon._connections import V1ConnectionResource
+from polyaxon._env_vars.keys import (
+    ENV_KEYS_API_VERSION,
+    ENV_KEYS_AUTH_TOKEN,
+    ENV_KEYS_AUTHENTICATION_TYPE,
+    ENV_KEYS_HAS_PROCESS_SIDECAR,
+    ENV_KEYS_HEADER,
+    ENV_KEYS_HEADER_SERVICE,
+    ENV_KEYS_HOST,
+    ENV_KEYS_IS_MANAGED,
+    ENV_KEYS_K8S_NAMESPACE,
+    ENV_KEYS_K8S_NODE_NAME,
+    ENV_KEYS_K8S_POD_ID,
+    ENV_KEYS_LOG_LEVEL,
+    ENV_KEYS_SECRET_INTERNAL_TOKEN,
+    ENV_KEYS_SECRET_KEY,
+)
+from polyaxon._local_process import process_types
+from polyaxon._runner.converter import BaseConverter
+from polyaxon._services.headers import PolyaxonServiceHeaders
+from polyaxon.api import VERSION_V1
+from polyaxon.exceptions import PolyaxonConverterError
+
+
+class EnvMixin(BaseConverter):
+    @staticmethod
+    def _get_env_var(name: str, value: Any) -> process_types.V1EnvVar:
+        if not isinstance(value, str):
+            try:
+                value = orjson_dumps(value)
+            except (ValueError, TypeError) as e:
+                raise PolyaxonConverterError(e)
+
+        return process_types.V1EnvVar(__root__=(name, value))
+
+    @staticmethod
+    def _get_from_json_resource(
+        resource: V1ConnectionResource,
+    ) -> List[process_types.V1EnvVar]:
+        if not resource or resource.items or resource.mount_path:
+            return []
+
+        secret = os.environ.get(resource.name)
+        if not secret:
+            return []
+        try:
+            secret_value = orjson_loads(secret)
+        except Exception as e:
+            raise PolyaxonConverterError from e
+
+        return [process_types.V1EnvVar(__root__=k) for k in (secret_value.items())]
+
+    @classmethod
+    def _get_env_from_json_resources(
+        cls,
+        resources: Iterable[V1ConnectionResource],
+    ) -> List[process_types.V1EnvVar]:
+        resources = resources or []
+        results = []
+        for resource in resources:
+            results += cls._get_from_json_resource(resource=resource)
+        return [r for r in results if r]
+
+    @staticmethod
+    def _get_item_from_json_resource(
+        key: str, resource_ref_name: str
+    ) -> Optional[process_types.V1EnvVar]:
+        secret = os.environ.get(resource_ref_name)
+        if not secret:
+            return None
+        try:
+            secret_value = orjson_loads(secret)
+        except orjson.JSONDecodeError:
+            return process_types.V1EnvVar(__root__=(key, secret))
+
+        value = secret_value.get(key)
+        return process_types.V1EnvVar(__root__=(key, value))
+
+    @classmethod
+    def _get_items_from_json_resource(
+        cls,
+        resource: V1ConnectionResource,
+    ) -> List[process_types.V1EnvVar]:
+        items_from = []
+        if not resource or not resource.items:
+            return items_from
+
+        try:
+            secret_value = orjson_loads(resource.name)
+        except orjson.JSONDecodeError as e:
+            return items_from
+
+        for item in resource.items:
+            value = secret_value.get(item)
+            if value:
+                items_from.append(process_types.V1EnvVar(__root__=(item, value)))
+        return items_from
+
+    @classmethod
+    def _get_env_vars_from_resources(
+        cls,
+        resources: Iterable[V1ConnectionResource],
+    ) -> List[process_types.V1EnvVar]:
+        resources = resources or []
+        env_vars = []
+        for secret in resources:
+            env_vars += cls._get_items_from_json_resource(resource=secret)
+        return env_vars
+
+    @classmethod
+    def _get_env_vars_from_k8s_resources(
+        cls,
+        secrets: Iterable[V1ConnectionResource],
+        config_maps: Iterable[V1ConnectionResource],
+    ) -> List[process_types.V1EnvVar]:
+        resources = to_list(secrets, check_none=True) + to_list(
+            config_maps, check_none=True
+        )
+        return cls._get_env_vars_from_resources(resources=resources)
+
+    @classmethod
+    def _get_additional_env_vars(cls) -> List[process_types.V1EnvVar]:
+        return [cls._get_env_var(name=ENV_KEYS_HAS_PROCESS_SIDECAR, value=True)]
+
+    @classmethod
+    def _get_base_env_vars(
+        cls,
+        namespace: str,
+        resource_name: str,
+        use_proxy_env_vars_use_in_ops: bool,
+        log_level: Optional[str] = None,
+    ) -> List[process_types.V1EnvVar]:
+        env = [
+            cls._get_env_var(name=ENV_KEYS_K8S_NODE_NAME, value="process-agent"),
+            cls._get_env_var(name=ENV_KEYS_K8S_NAMESPACE, value=namespace),
+            cls._get_env_var(name=ENV_KEYS_K8S_POD_ID, value=resource_name),
+        ]
+        if log_level:
+            env.append(cls._get_env_var(name=ENV_KEYS_LOG_LEVEL, value=log_level))
+        env += cls._get_proxy_env_vars(use_proxy_env_vars_use_in_ops)
+        return env
+
+    def _get_service_env_vars(
+        self,
+        service_header: str,
+        header: Optional[str] = PolyaxonServiceHeaders.SERVICE,
+        include_secret_key: bool = False,
+        include_internal_token: bool = False,
+        include_agent_token: bool = False,
+        authentication_type: Optional[str] = None,
+        external_host: bool = False,
+        log_level: Optional[str] = None,
+        polyaxon_default_secret_ref: Optional[str] = None,
+        polyaxon_agent_secret_ref: Optional[str] = None,
+        api_version: Optional[str] = None,
+        use_proxy_env_vars_use_in_ops: Optional[bool] = None,
+    ) -> List[process_types.V1EnvVar]:
+        api_host = self.get_api_host(external_host)
+        polyaxon_default_secret_ref = (
+            polyaxon_default_secret_ref or settings.AGENT_CONFIG.app_secret_name
+        )
+        polyaxon_agent_secret_ref = (
+            polyaxon_agent_secret_ref or settings.AGENT_CONFIG.agent_secret_name
+        )
+        # TODO: Add support for local CLI auth
+        use_proxy_env_vars_use_in_ops = (
+            use_proxy_env_vars_use_in_ops
+            if use_proxy_env_vars_use_in_ops is not None
+            else settings.AGENT_CONFIG.use_proxy_env_vars_use_in_ops
+        )
+        api_version = api_version or VERSION_V1
+        env_vars = self._get_base_env_vars(
+            namespace=self.namespace,
+            resource_name=self.resource_name,
+            use_proxy_env_vars_use_in_ops=use_proxy_env_vars_use_in_ops,
+        ) + [
+            self._get_env_var(name=ENV_KEYS_HOST, value=api_host),
+            self._get_env_var(name=ENV_KEYS_IS_MANAGED, value=True),
+            self._get_env_var(name=ENV_KEYS_API_VERSION, value=api_version),
+            self._get_run_instance_env_var(self.run_instance),
+        ]
+        if log_level:
+            env_vars.append(self._get_env_var(name=ENV_KEYS_LOG_LEVEL, value=log_level))
+        if header:
+            env_vars.append(
+                self._get_env_var(
+                    name=ENV_KEYS_HEADER,
+                    value=PolyaxonServiceHeaders.get_header(header),
+                )
+            )
+        if service_header:
+            env_vars += to_list(
+                self._get_env_var(
+                    name=ENV_KEYS_HEADER_SERVICE, value=get_enum_value(service_header)
+                ),
+            )
+        if include_secret_key:
+            env_vars += to_list(
+                self._get_item_from_json_resource(
+                    key=ENV_KEYS_SECRET_KEY,
+                    resource_ref_name=polyaxon_default_secret_ref,
+                ),
+                check_none=True,
+            )
+        internal = False
+        if include_internal_token and polyaxon_default_secret_ref:
+            internal = True
+            env_vars += to_list(
+                self._get_item_from_json_resource(
+                    key=ENV_KEYS_SECRET_INTERNAL_TOKEN,
+                    resource_ref_name=polyaxon_default_secret_ref,
+                ),
+                check_none=True,
+            )
+        if include_agent_token:
+            if internal:
+                raise PolyaxonConverterError(
+                    "A service cannot have internal token and agent token."
+                )
+            if polyaxon_agent_secret_ref:
+                env_vars += to_list(
+                    self._get_item_from_json_resource(
+                        key=ENV_KEYS_AUTH_TOKEN,
+                        resource_ref_name=polyaxon_agent_secret_ref,
+                    ),
+                    check_none=True,
+                )
+            elif settings.CLIENT_CONFIG.token:
+                env_vars += to_list(
+                    self._get_env_var(
+                        name=ENV_KEYS_AUTH_TOKEN, value=settings.CLIENT_CONFIG.token
+                    ),
+                    check_none=True,
+                )
+        if authentication_type:
+            env_vars += to_list(
+                self._get_env_var(
+                    name=ENV_KEYS_AUTHENTICATION_TYPE, value=authentication_type
+                ),
+                check_none=True,
+            )
+        return env_vars