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

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)],
+        )

polyaxon/_local_process/converter/base/main.py

@@ -0,0 +1,74 @@
+from typing import Dict, Iterable, List, Optional
+
+from clipped.utils.lists import to_list
+
+from polyaxon._connections import V1Connection, V1ConnectionResource
+from polyaxon._flow import V1Init, V1Plugins
+from polyaxon._local_process import process_types
+from polyaxon._runner.converter import BaseConverter as _BaseConverter
+from polyaxon.exceptions import PolyaxonConverterError
+
+
+class MainConverter(_BaseConverter):
+    def _get_main_container(
+        self,
+        container_id: str,
+        main_container: process_types.V1Container,
+        plugins: V1Plugins,
+        artifacts_store: Optional[V1Connection],
+        init: Optional[List[V1Init]],
+        connections: Optional[List[str]],
+        connection_by_names: Dict[str, V1Connection],
+        secrets: Optional[Iterable[V1ConnectionResource]],
+        config_maps: Optional[Iterable[V1ConnectionResource]],
+        run_path: Optional[str],
+        kv_env_vars: List[List] = None,
+        ports: List[int] = None,
+    ) -> process_types.V1Container:
+        connections = connections or []
+        connection_by_names = connection_by_names or {}
+        secrets = secrets or []
+        config_maps = config_maps or []
+
+        if artifacts_store and not run_path:
+            raise PolyaxonConverterError("Run path is required for main container.")
+
+        if artifacts_store and (
+            not plugins.collect_artifacts or plugins.mount_artifacts_store
+        ):
+            if artifacts_store.name not in connection_by_names:
+                connection_by_names[artifacts_store.name] = artifacts_store
+            if artifacts_store.name not in connections:
+                connections.append(artifacts_store.name)
+
+        requested_connections = [connection_by_names[c] for c in connections]
+        requested_config_maps = V1Connection.get_requested_resources(
+            resources=config_maps,
+            connections=requested_connections,
+            resource_key="config_map",
+        )
+        requested_secrets = V1Connection.get_requested_resources(
+            resources=secrets, connections=requested_connections, resource_key="secret"
+        )
+
+        # Env vars
+        env = self._get_main_env_vars(
+            plugins=plugins,
+            kv_env_vars=kv_env_vars,
+            artifacts_store_name=artifacts_store.name if artifacts_store else None,
+            connections=requested_connections,
+            secrets=requested_secrets,
+            config_maps=requested_config_maps,
+        )
+
+        # Env from
+        resources = to_list(requested_secrets, check_none=True) + to_list(
+            requested_config_maps, check_none=True
+        )
+        env += self._get_env_from_json_resources(resources=resources)
+
+        return self._patch_container(
+            container=main_container,
+            name=container_id,
+            env=env,
+        )

polyaxon/_local_process/converter/base/mounts.py

@@ -0,0 +1,82 @@
+from typing import List, Optional
+
+from polyaxon._connections import V1Connection, V1ConnectionResource
+from polyaxon._contexts import paths as ctx_paths
+from polyaxon._local_process import process_types
+from polyaxon._runner.converter import BaseConverter
+
+
+class MountsMixin(BaseConverter):
+    @classmethod
+    def _get_mount_from_store(
+        cls,
+        store: V1Connection,
+    ):
+        pass
+
+    @classmethod
+    def _get_mount_from_resource(
+        cls,
+        resource: V1ConnectionResource,
+    ):
+        pass
+
+    @classmethod
+    def _get_volume(
+        cls,
+        mount_path: str,
+        host_path: Optional[str] = None,
+        read_only: Optional[bool] = None,
+    ):
+        pass
+
+    @classmethod
+    def _get_docker_context_mount(cls):
+        pass
+
+    @classmethod
+    def _get_auth_context_mount(
+        cls,
+        read_only: Optional[bool] = None,
+        run_path: Optional[str] = None,
+    ):
+        pass
+
+    @classmethod
+    def _get_artifacts_context_mount(
+        cls,
+        read_only: bool = False,
+        run_path: Optional[str] = None,
+    ):
+        pass
+
+    @classmethod
+    def _get_connections_context_mount(
+        cls,
+        name: str,
+        mount_path: str,
+        run_path: str,
+    ):
+        pass
+
+    @classmethod
+    def _get_shm_context_mount(cls):
+        """
+        Mount a tmpfs volume to /dev/shm.
+        This will set /dev/shm size to half of the RAM of node.
+        By default, /dev/shm is very small, only 64MB.
+        Some experiments will fail due to lack of share memory,
+        such as some experiments running on Pytorch.
+        """
+        pass
+
+    @classmethod
+    def _get_mounts(
+        cls,
+        use_auth_context: bool,
+        use_docker_context: bool,
+        use_shm_context: bool,
+        use_artifacts_context: bool,
+        run_path: Optional[str] = None,
+    ) -> List:
+        return []

polyaxon/_local_process/converter/converters/__init__.py

@@ -0,0 +1,8 @@
+from polyaxon._flow import V1RunKind
+from polyaxon._local_process.converter.converters.job import JobConverter
+from polyaxon._local_process.converter.converters.service import ServiceConverter
+
+CONVERTERS = {
+    V1RunKind.JOB: JobConverter,
+    V1RunKind.SERVICE: ServiceConverter,
+}

polyaxon/_local_process/converter/converters/job.py

@@ -0,0 +1,40 @@
+from typing import Dict, Iterable, List, Optional
+
+from polyaxon._connections import V1Connection, V1ConnectionResource
+from polyaxon._flow import V1CompiledOperation, V1Job, V1Plugins
+from polyaxon._k8s.converter.mixins import JobMixin
+from polyaxon._local_process import process_types
+from polyaxon._local_process.converter.base import BaseConverter
+
+
+class JobConverter(JobMixin, BaseConverter):
+    def get_resource(
+        self,
+        compiled_operation: V1CompiledOperation,
+        artifacts_store: V1Connection,
+        connection_by_names: Dict[str, V1Connection],
+        secrets: Optional[Iterable[V1ConnectionResource]],
+        config_maps: Optional[Iterable[V1ConnectionResource]],
+        default_sa: Optional[str] = None,
+        default_auth: bool = False,
+    ) -> List[process_types.V1Container]:
+        job = compiled_operation.run  # type: V1Job
+        plugins = V1Plugins.get_or_create(
+            config=compiled_operation.plugins, auth=default_auth
+        )
+        kv_env_vars = compiled_operation.get_env_io()
+        return self.get_replica_resource(
+            environment=job.environment,
+            plugins=plugins,
+            volumes=job.volumes,
+            init=job.init,
+            sidecars=job.sidecars,
+            container=job.container,
+            artifacts_store=artifacts_store,
+            connections=job.connections,
+            connection_by_names=connection_by_names,
+            secrets=secrets,
+            config_maps=config_maps,
+            kv_env_vars=kv_env_vars,
+            default_sa=default_sa,
+        )

polyaxon/_local_process/converter/converters/service.py

@@ -0,0 +1,41 @@
+from typing import Dict, Iterable, List, Optional
+
+from polyaxon._connections import V1Connection, V1ConnectionResource
+from polyaxon._flow import V1CompiledOperation, V1Plugins, V1Service
+from polyaxon._local_process import process_types
+from polyaxon._local_process.converter.base import BaseConverter
+from polyaxon._local_process.converter.mixins import ServiceMixin
+
+
+class ServiceConverter(ServiceMixin, BaseConverter):
+    def get_resource(
+        self,
+        compiled_operation: V1CompiledOperation,
+        artifacts_store: V1Connection,
+        connection_by_names: Dict[str, V1Connection],
+        secrets: Optional[Iterable[V1ConnectionResource]],
+        config_maps: Optional[Iterable[V1ConnectionResource]],
+        default_sa: Optional[str] = None,
+        default_auth: bool = False,
+    ) -> List[process_types.V1Container]:
+        service = compiled_operation.run  # type: V1Service
+        plugins = V1Plugins.get_or_create(
+            config=compiled_operation.plugins, auth=default_auth
+        )
+        kv_env_vars = compiled_operation.get_env_io()
+        return self.get_replica_resource(
+            plugins=plugins,
+            environment=service.environment,
+            volumes=service.volumes,
+            init=service.init,
+            sidecars=service.sidecars,
+            container=service.container,
+            artifacts_store=artifacts_store,
+            connections=service.connections,
+            connection_by_names=connection_by_names,
+            secrets=secrets,
+            config_maps=config_maps,
+            kv_env_vars=kv_env_vars,
+            default_sa=default_sa,
+            ports=service.ports,
+        )

polyaxon/_local_process/converter/mixins.py

@@ -0,0 +1,38 @@
+from typing import Dict
+
+from polyaxon._containers.names import MAIN_JOB_CONTAINER
+from polyaxon._flow import V1RunKind
+
+
+class JobMixin:
+    K8S_ANNOTATIONS_KIND = V1RunKind.JOB
+    MAIN_CONTAINER_ID = MAIN_JOB_CONTAINER
+
+
+class NotifierMixin:
+    K8S_ANNOTATIONS_KIND = V1RunKind.NOTIFIER
+    MAIN_CONTAINER_ID = MAIN_JOB_CONTAINER
+
+
+class CleanerMixin:
+    K8S_ANNOTATIONS_KIND = V1RunKind.CLEANER
+    MAIN_CONTAINER_ID = MAIN_JOB_CONTAINER
+
+
+class TunerMixin:
+    K8S_ANNOTATIONS_KIND = V1RunKind.TUNER
+    MAIN_CONTAINER_ID = MAIN_JOB_CONTAINER
+
+
+class ServiceMixin:
+    K8S_ANNOTATIONS_KIND = V1RunKind.SERVICE
+    MAIN_CONTAINER_ID = MAIN_JOB_CONTAINER
+
+
+MIXIN_MAPPING: Dict = {
+    V1RunKind.JOB: JobMixin,
+    V1RunKind.NOTIFIER: NotifierMixin,
+    V1RunKind.CLEANER: CleanerMixin,
+    V1RunKind.TUNER: TunerMixin,
+    V1RunKind.SERVICE: ServiceMixin,
+}