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

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.make((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.make(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.make((key, secret))
+
+        value = secret_value.get(key)
+        return process_types.V1EnvVar.make((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.make((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

polyaxon/_local_process/converter/base/init.py

@@ -0,0 +1,414 @@
+from typing import Any, List, Optional, Tuple, Union
+
+from clipped.utils.enums import get_enum_value
+from clipped.utils.lists import to_list
+
+from polyaxon._auxiliaries import V1PolyaxonInitContainer
+from polyaxon._connections import V1Connection, V1ConnectionKind
+from polyaxon._constants.globals import DEFAULT
+from polyaxon._containers.names import (
+    INIT_ARTIFACTS_CONTAINER_PREFIX,
+    INIT_AUTH_CONTAINER,
+    INIT_CUSTOM_CONTAINER_PREFIX,
+    INIT_DOCKERFILE_CONTAINER_PREFIX,
+    INIT_FILE_CONTAINER_PREFIX,
+    INIT_GIT_CONTAINER_PREFIX,
+    INIT_TENSORBOARD_CONTAINER_PREFIX,
+    generate_container_name,
+)
+from polyaxon._contexts import paths as ctx_paths
+from polyaxon._env_vars.keys import ENV_KEYS_SSH_PATH
+from polyaxon._flow import V1Plugins
+from polyaxon._local_process import process_types
+from polyaxon._runner.converter import BaseConverter as _BaseConverter
+from polyaxon._runner.converter.init.artifacts import init_artifact_context_args
+from polyaxon._runner.converter.init.file import FILE_INIT_COMMAND, get_file_init_args
+from polyaxon._runner.converter.init.git import REPO_INIT_COMMAND, get_repo_context_args
+from polyaxon._runner.converter.init.store import get_volume_args
+from polyaxon._runner.converter.init.tensorboard import (
+    TENSORBOARD_INIT_COMMAND,
+    get_tensorboard_args,
+)
+from polyaxon._schemas.types import (
+    V1ArtifactsType,
+    V1DockerfileType,
+    V1FileType,
+    V1TensorboardType,
+)
+from polyaxon.exceptions import PolyaxonConverterError
+
+
+class InitConverter(_BaseConverter):
+    @classmethod
+    def _get_base_store_container(
+        cls,
+        container: Optional[process_types.V1Container],
+        container_name: str,
+        polyaxon_init: V1PolyaxonInitContainer,
+        store: V1Connection,
+        env: List[process_types.V1EnvVar],
+        env_from: List[Any],
+        args: List[str],
+        command: Optional[List[str]] = None,
+        **kwargs
+    ) -> Optional[process_types.V1Container]:
+        env = env or []
+
+        # Artifact store needs to allow init the contexts as well, so the store is not required
+        if not store:
+            raise PolyaxonConverterError("Init store container requires a store")
+
+        if store.is_bucket:
+            secret = store.secret
+            env = env + to_list(
+                cls._get_items_from_json_resource(resource=secret), check_none=True
+            )
+            config_map = store.config_map
+            env = env + to_list(
+                cls._get_from_json_resource(resource=config_map),
+                check_none=True,
+            )
+        # Add connections catalog env vars information
+        connection_env = cls._get_connections_catalog_env_var(connections=[store])
+        if connection_env:
+            env.append(connection_env)
+
+        return cls._patch_container(
+            container=container,
+            name=container_name,
+            command=command or ["/bin/sh", "-c"],
+            args=args,
+            env=env,
+        )
+
+    @classmethod
+    def _get_custom_init_container(
+        cls,
+        connection: V1Connection,
+        plugins: V1Plugins,
+        container: Optional[process_types.V1Container],
+        run_path: str,
+        env: List[process_types.V1EnvVar] = None,
+        mount_path: Optional[str] = None,
+    ) -> process_types.V1Container:
+        if not connection:
+            raise PolyaxonConverterError(
+                "A connection is required to create a repo context."
+            )
+
+        env = to_list(env, check_none=True)
+        secret = connection.secret
+        if secret:
+            env += to_list(
+                cls._get_from_json_resource(resource=secret), check_none=True
+            )
+
+        # Add connections catalog env vars information
+        connection_env = cls._get_connections_catalog_env_var(connections=[connection])
+        if connection_env:
+            env.append(connection_env)
+        connection_env = cls._get_connection_env_var(connection=connection)
+        if connection_env:
+            env.append(connection_env)
+        config_map = connection.config_map
+        if config_map:
+            env += to_list(
+                cls._get_from_json_resource(resource=config_map),
+                check_none=True,
+            )
+        container_name = container.name or generate_container_name(
+            INIT_CUSTOM_CONTAINER_PREFIX, connection.name
+        )
+        return cls._patch_container(
+            container=container,
+            name=container_name,
+            env=env,
+        )
+
+    @classmethod
+    def _get_dockerfile_init_container(
+        cls,
+        polyaxon_init: V1PolyaxonInitContainer,
+        dockerfile_args: V1DockerfileType,
+        plugins: V1Plugins,
+        run_path: str,
+        run_instance: str,
+        container: Optional[process_types.V1Container] = None,
+        env: List[process_types.V1EnvVar] = None,
+        mount_path: Optional[str] = None,
+    ) -> process_types.V1Container:
+        env = to_list(env, check_none=True)
+        env = env + [cls._get_run_instance_env_var(run_instance)]
+
+        container_name = generate_container_name(INIT_DOCKERFILE_CONTAINER_PREFIX)
+        if not container:
+            container = cls._new_container(name=container_name)
+
+        mount_path = mount_path or ctx_paths.CONTEXT_MOUNT_ARTIFACTS
+
+        return cls._patch_container(
+            container=container,
+            name=container_name,
+            image=polyaxon_init.get_image(),
+            image_pull_policy=polyaxon_init.image_pull_policy,
+            command=["polyaxon", "docker", "generate"],
+            args=[
+                "--build-context={}".format(dockerfile_args.to_json()),
+                "--destination={}".format(mount_path),
+                "--copy-path={}".format(
+                    ctx_paths.CONTEXT_MOUNT_RUN_OUTPUTS_FORMAT.format(run_path)
+                ),
+                "--track",
+            ],
+            env=env,
+        )
+
+    @classmethod
+    def _get_file_init_container(
+        cls,
+        polyaxon_init: V1PolyaxonInitContainer,
+        file_args: V1FileType,
+        plugins: V1Plugins,
+        run_path: str,
+        run_instance: str,
+        container: Optional[process_types.V1Container] = None,
+        env: List[process_types.V1EnvVar] = None,
+        mount_path: Optional[str] = None,
+    ) -> process_types.V1Container:
+        env = to_list(env, check_none=True)
+        env = env + [cls._get_run_instance_env_var(run_instance)]
+
+        container_name = generate_container_name(INIT_FILE_CONTAINER_PREFIX)
+        if not container:
+            container = cls._new_container(name=container_name)
+
+        mount_path = mount_path or ctx_paths.CONTEXT_MOUNT_ARTIFACTS
+
+        file_args.filename = file_args.filename or "file"
+        return cls._patch_container(
+            container=container,
+            name=container_name,
+            image=polyaxon_init.get_image(),
+            image_pull_policy=polyaxon_init.image_pull_policy,
+            command=FILE_INIT_COMMAND,
+            args=get_file_init_args(
+                file_args=file_args, run_path=run_path, mount_path=mount_path
+            ),
+            env=env,
+        )
+
+    @classmethod
+    def _get_git_init_container(
+        cls,
+        polyaxon_init: V1PolyaxonInitContainer,
+        connection: V1Connection,
+        plugins: V1Plugins,
+        run_path: str,
+        container: Optional[process_types.V1Container] = None,
+        env: List[process_types.V1EnvVar] = None,
+        mount_path: Optional[str] = None,
+        track: bool = False,
+    ) -> process_types.V1Container:
+        if not connection:
+            raise PolyaxonConverterError(
+                "A connection is required to create a repo context."
+            )
+        container_name = generate_container_name(
+            INIT_GIT_CONTAINER_PREFIX, connection.name
+        )
+        if not container:
+            container = cls._new_container(name=container_name)
+
+        mount_path = mount_path or ctx_paths.CONTEXT_MOUNT_ARTIFACTS
+
+        env = to_list(env, check_none=True)
+        env_from = []
+        secret = connection.secret
+        if secret:
+            env += to_list(
+                cls._get_from_json_resource(resource=secret), check_none=True
+            )
+            env_from += to_list(
+                cls._get_env_from_secret(secret=secret), check_none=True
+            )
+
+        # Add connections catalog env vars information
+        env += to_list(
+            cls._get_connections_catalog_env_var(connections=[connection]),
+            check_none=True,
+        )
+        env += to_list(
+            cls._get_connection_env_var(connection=connection), check_none=True
+        )
+        # Add special handling to auto-inject ssh mount path
+        if connection.kind == V1ConnectionKind.SSH and secret.mount_path:
+            env += [cls._get_env_var(ENV_KEYS_SSH_PATH, secret.mount_path)]
+        config_map = connection.config_map
+        if config_map:
+            env += to_list(
+                cls._get_from_json_resource(resource=config_map),
+                check_none=True,
+            )
+            env_from += to_list(
+                cls._get_env_from_config_map(config_map=config_map),
+                check_none=True,
+            )
+        args = get_repo_context_args(
+            name=connection.name,
+            # Handle the case of custom connection
+            url=getattr(connection.schema_, "url", None),
+            revision=getattr(connection.schema_, "revision", None),
+            flags=getattr(connection.schema_, "flags", None),
+            mount_path=mount_path,
+            connection=connection.name if track else None,
+        )
+        return cls._patch_container(
+            container=container,
+            name=container_name,
+            image=polyaxon_init.get_image(),
+            image_pull_policy=polyaxon_init.image_pull_policy,
+            command=REPO_INIT_COMMAND,
+            args=args,
+            env=env,
+            env_from=env_from,
+        )
+
+    @classmethod
+    def _get_store_init_container(
+        cls,
+        polyaxon_init: V1PolyaxonInitContainer,
+        connection: V1Connection,
+        artifacts: V1ArtifactsType,
+        paths: Union[List[str], List[Tuple[str, str]]],
+        run_path: str,
+        container: Optional[process_types.V1Container] = None,
+        env: List[process_types.V1EnvVar] = None,
+        mount_path: Optional[str] = None,
+        is_default_artifacts_store: bool = False,
+    ) -> process_types.V1Container:
+        container_name = generate_container_name(
+            INIT_ARTIFACTS_CONTAINER_PREFIX, connection.name
+        )
+        if not container:
+            container = cls._new_container(name=container_name)
+
+        mount_path = mount_path or (
+            ctx_paths.CONTEXT_MOUNT_ARTIFACTS
+            if is_default_artifacts_store
+            else ctx_paths.CONTEXT_MOUNT_ARTIFACTS_FORMAT.format(connection.name)
+        )
+
+        return cls._get_base_store_container(
+            container=container,
+            container_name=container_name,
+            polyaxon_init=polyaxon_init,
+            store=connection,
+            env=env,
+            env_from=[],
+            args=[
+                get_volume_args(
+                    store=connection,
+                    mount_path=mount_path,
+                    artifacts=artifacts,
+                    paths=paths,
+                )
+            ],
+        )
+
+    @classmethod
+    def _get_tensorboard_init_container(
+        cls,
+        polyaxon_init: V1PolyaxonInitContainer,
+        artifacts_store: V1Connection,
+        tb_args: V1TensorboardType,
+        plugins: V1Plugins,
+        run_path: str,
+        run_instance: str,
+        container: Optional[process_types.V1Container] = None,
+        env: List[process_types.V1EnvVar] = None,
+        mount_path: Optional[str] = None,
+    ) -> process_types.V1Container:
+        env = to_list(env, check_none=True)
+        env = env + [cls._get_run_instance_env_var(run_instance)]
+
+        container_name = generate_container_name(INIT_TENSORBOARD_CONTAINER_PREFIX)
+        if not container:
+            container = cls._new_container(name=container_name)
+
+        mount_path = mount_path or ctx_paths.CONTEXT_MOUNT_ARTIFACTS
+
+        args = get_tensorboard_args(
+            tb_args=tb_args,
+            context_from=artifacts_store.store_path,
+            context_to=mount_path,
+            connection_kind=get_enum_value(artifacts_store.kind),
+        )
+
+        return cls._get_base_store_container(
+            container=container,
+            container_name=container_name,
+            polyaxon_init=polyaxon_init,
+            store=artifacts_store,
+            command=TENSORBOARD_INIT_COMMAND,
+            args=args,
+            env=env,
+            env_from=[],
+        )
+
+    @classmethod
+    def _get_auth_context_init_container(
+        cls,
+        polyaxon_init: V1PolyaxonInitContainer,
+        run_path: str,
+        env: Optional[List[process_types.V1EnvVar]] = None,
+    ) -> process_types.V1Container:
+        env = to_list(env, check_none=True)
+        container = process_types.V1Container(
+            name=INIT_AUTH_CONTAINER,
+            image=polyaxon_init.get_image(),
+            command=["polyaxon", "initializer", "auth"],
+            env=env,
+        )
+        return cls._patch_container(container)
+
+    @classmethod
+    def _get_artifacts_path_init_container(
+        cls,
+        polyaxon_init: V1PolyaxonInitContainer,
+        artifacts_store: V1Connection,
+        run_path: str,
+        auto_resume: bool,
+        env: Optional[List[process_types.V1EnvVar]] = None,
+    ) -> process_types.V1Container:
+        if not artifacts_store:
+            raise PolyaxonConverterError("Init artifacts container requires a store.")
+
+        env = to_list(env, check_none=True)
+        init_args = init_artifact_context_args(run_path=run_path)
+        if auto_resume:
+            init_args.append(
+                get_volume_args(
+                    store=artifacts_store,
+                    mount_path=ctx_paths.CONTEXT_MOUNT_ARTIFACTS,
+                    artifacts=V1ArtifactsType(dirs=[run_path]),
+                    paths=None,
+                    sync_fw=True,
+                )
+            )
+
+        container_name = generate_container_name(
+            INIT_ARTIFACTS_CONTAINER_PREFIX, DEFAULT, False
+        )
+        container = cls._new_container(name=container_name)
+
+        return cls._get_base_store_container(
+            container_name=container_name,
+            container=container,
+            polyaxon_init=polyaxon_init,
+            store=artifacts_store,
+            env=env,
+            env_from=[],
+            # If we are dealing with a volume we need to make sure the path exists for the user
+            # We also clean the path if this is not a resume run
+            args=[" ".join(init_args)],
+        )