Wire sandbox plugin into K8s services

 * Copy plx-exec and the sandbox bootstrap script through the init tools container, wrap sandbox service commands, inject the derived sandbox token, and expose port 9090 on the pod and service spec.

Author: md-mq

cli/polyaxon/_env_vars/keys.py

@@ -106,6 +106,7 @@
 ENV_KEYS_DASKCLUSTER_ENABLED = "POLYAXON_DASKCLUSTER_ENABLED"
 
 # Sandbox
+ENV_KEYS_SANDBOX_TOKEN = "POLYAXON_SANDBOX_TOKEN"
 ENV_KEYS_SANDBOX_PORT = "POLYAXON_SANDBOX_PORT"
 ENV_KEYS_SANDBOX_HOST = "POLYAXON_SANDBOX_HOST"
 ENV_KEYS_SANDBOX_DEBUG = "POLYAXON_SANDBOX_DEBUG"

cli/polyaxon/_flow/plugins/__init__.py

@@ -159,6 +159,12 @@ class V1Plugins(BaseSchemaModel):
     This plugin enables the sandbox daemon for programmatic exec, filesystem,
     and PTY access from SDKs, agents, and UIs.
 
+    Sandbox is supported for service runs only. When enabled, Polyaxon wraps the
+    main container command with `/opt/polyaxon/bin/bootstrap-sandbox.sh`. If no
+    command is provided, `plx-exec` runs as PID 1. If a workload is needed, set
+    an explicit `command`; args without command are not supported in v0. The
+    user image must provide `/bin/sh`.
+
     To enable this plugin:
 
     ```yaml

cli/polyaxon/_k8s/converter/base/init.py

@@ -515,13 +515,31 @@ def _get_auth_context_init_container(
     def _get_tools_init_container(
         cls,
         polyaxon_init: V1PolyaxonInitContainer,
+        use_tmux: bool = True,
+        use_sandbox: bool = False,
     ) -> k8s_schemas.V1Container:
+        if not use_tmux and not use_sandbox:
+            raise PolyaxonConverterError("Init tools container requires a tool.")
+
+        copy_commands = []
+        if use_tmux:
+            copy_commands.append("cp /usr/bin/tmux /opt/polyaxon/bin/tmux")
+        if use_sandbox:
+            copy_commands += [
+                "cp /usr/bin/plx-exec /opt/polyaxon/bin/plx-exec",
+                "cp /usr/bin/bootstrap-sandbox.sh "
+                "/opt/polyaxon/bin/bootstrap-sandbox.sh",
+            ]
+        command = ["sh", "-c", " && ".join(copy_commands)]
+        if use_tmux and not use_sandbox:
+            command = ["cp", "/usr/bin/tmux", "/opt/polyaxon/bin/tmux"]
+
         return cls._patch_container(
             container=k8s_schemas.V1Container(
                 name=INIT_TOOLS_CONTAINER,
                 image=polyaxon_init.get_image(),
                 image_pull_policy=polyaxon_init.image_pull_policy,
-                command=["cp", "/usr/bin/tmux", "/opt/polyaxon/bin/tmux"],
+                command=command,
                 resources=polyaxon_init.get_resources(),
                 volume_mounts=[cls._get_tools_bin_context_mount(read_only=False)],
             )

cli/polyaxon/_k8s/converter/base/main.py

@@ -3,12 +3,31 @@
 from clipped.utils.lists import to_list
 
 from polyaxon._connections import V1Connection, V1ConnectionResource
+from polyaxon._env_vars.keys import ENV_KEYS_SANDBOX_TOKEN
 from polyaxon._flow import V1Init, V1Plugins
 from polyaxon._k8s import k8s_schemas
 from polyaxon._runner.converter import BaseConverter as _BaseConverter
+from polyaxon._sandbox.auth import derive_sandbox_token_from_env
+from polyaxon._sandbox.constants import SANDBOX_BOOTSTRAP_PATH, SANDBOX_PORT
 from polyaxon.exceptions import PolyaxonConverterError
 
 
+def _has_container_port(ports, port: int) -> bool:
+    for container_port in to_list(ports, check_none=True):
+        value = getattr(container_port, "container_port", None)
+        if value is None and isinstance(container_port, dict):
+            value = container_port.get("containerPort") or container_port.get(
+                "container_port"
+            )
+        if str(value) == str(port):
+            return True
+    return False
+
+
+def _has_port(ports, port: int) -> bool:
+    return any(str(p) == str(port) for p in to_list(ports, check_none=True))
+
+
 class MainConverter(_BaseConverter):
     def _get_main_container(
         self,
@@ -33,6 +52,21 @@ def _get_main_container(
         if artifacts_store and not run_path:
             raise PolyaxonConverterError("Run path is required for main container.")
 
+        if plugins and plugins.sandbox:
+            if not main_container:
+                raise PolyaxonConverterError(
+                    "plugins.sandbox requires a main container."
+                )
+            if main_container.args and not main_container.command:
+                raise PolyaxonConverterError(
+                    "plugins.sandbox does not support args without command."
+                )
+            user_argv = to_list(main_container.command, check_none=True) + to_list(
+                main_container.args, check_none=True
+            )
+            main_container.command = [SANDBOX_BOOTSTRAP_PATH]
+            main_container.args = user_argv
+
         if artifacts_store and (
             not plugins.collect_artifacts or plugins.mount_artifacts_store
         ):
@@ -59,6 +93,7 @@ def _get_main_container(
                 use_docker_context=plugins.docker,
                 use_shm_context=plugins.shm,
                 use_tmux_context=plugins.tmux,
+                use_sandbox_context=plugins.sandbox,
                 run_path=run_path,
             )
             if plugins
@@ -82,16 +117,38 @@ def _get_main_container(
             secrets=requested_secrets,
             config_maps=requested_config_maps,
         )
+        if plugins and plugins.sandbox:
+            env.append(
+                self._get_env_var(
+                    name=ENV_KEYS_SANDBOX_TOKEN,
+                    value=derive_sandbox_token_from_env(self.run_uuid),
+                )
+            )
         env += self._get_resources_env_vars(main_container.resources)
 
         # Env from
         env_from = self._get_env_from_k8s_resources(
             secrets=requested_secrets, config_maps=requested_config_maps
         )
 
+        ports = list(to_list(ports, check_none=True))
+        if plugins and plugins.sandbox:
+            ports = [
+                port
+                for port in ports
+                if not _has_container_port(main_container.ports, port)
+            ]
+        if (
+            plugins
+            and plugins.sandbox
+            and not _has_port(ports, SANDBOX_PORT)
+            and not _has_container_port(main_container.ports, SANDBOX_PORT)
+        ):
+            ports.append(SANDBOX_PORT)
+
         ports = [
             k8s_schemas.V1ContainerPort(container_port=port)
-            for port in to_list(ports, check_none=True)
+            for port in ports
         ]
 
         return self._patch_container(

cli/polyaxon/_k8s/converter/base/mounts.py

@@ -98,6 +98,7 @@ def _get_mounts(
         use_shm_context: bool,
         use_artifacts_context: bool,
         use_tmux_context: bool = False,
+        use_sandbox_context: bool = False,
         run_path: Optional[str] = None,
     ) -> List[k8s_schemas.V1VolumeMount]:
         mounts = []
@@ -113,7 +114,7 @@ def _get_mounts(
             mounts.append(cls._get_docker_context_mount())
         if use_shm_context:
             mounts.append(cls._get_shm_context_mount())
-        if use_tmux_context:
+        if use_tmux_context or use_sandbox_context:
             mounts.append(cls._get_tools_bin_context_mount(read_only=True))
 
         return mounts

cli/polyaxon/_k8s/converter/converters/service.py

@@ -1,13 +1,23 @@
 from typing import Dict, Iterable, Optional
 
+from clipped.utils.lists import to_list
+
 from polyaxon._connections import V1Connection, V1ConnectionResource
 from polyaxon._flow import V1CompiledOperation, V1Plugins
 from polyaxon._k8s.converter.base import BaseConverter
 from polyaxon._k8s.converter.mixins import ServiceMixin
 from polyaxon._k8s.custom_resources.service import get_service_custom_resource
+from polyaxon._sandbox.constants import SANDBOX_PORT
 
 
 class ServiceConverter(ServiceMixin, BaseConverter):
+    @staticmethod
+    def _get_service_ports(ports, plugins: V1Plugins):
+        ports = list(to_list(ports, check_none=True))
+        if plugins and plugins.sandbox and SANDBOX_PORT not in ports:
+            ports.append(SANDBOX_PORT)
+        return ports
+
     def get_resource(
         self,
         compiled_operation: V1CompiledOperation,
@@ -23,6 +33,7 @@ def get_resource(
             config=compiled_operation.plugins, auth=default_auth
         )
         kv_env_vars = compiled_operation.get_env_io()
+        ports = self._get_service_ports(service.ports, plugins)
         replica_spec = self.get_replica_resource(
             plugins=plugins,
             environment=service.environment,
@@ -37,7 +48,7 @@ def get_resource(
             config_maps=config_maps,
             kv_env_vars=kv_env_vars,
             default_sa=default_sa,
-            ports=service.ports,
+            ports=ports,
         )
         return get_service_custom_resource(
             namespace=self.namespace,
@@ -53,7 +64,7 @@ def get_resource(
             notifications=plugins.notifications,
             labels=replica_spec.labels,
             annotations=replica_spec.annotations,
-            ports=service.ports,
+            ports=ports,
             is_external=service.is_external,
             replicas=service.replicas,
         )

cli/polyaxon/_k8s/converter/pod/volumes.py

@@ -114,6 +114,6 @@ def add_volume_from_resource(resource: V1ConnectionResource, is_secret: bool):
         volumes.append(get_configs_context_volume())
     if plugins and plugins.docker:
         volumes.append(get_docker_context_volume())
-    if plugins and plugins.tmux:
+    if plugins and (plugins.tmux or plugins.sandbox):
         volumes.append(get_tools_bin_context_volume())
     return volumes

cli/polyaxon/_runner/converter/converter.py

@@ -501,6 +501,7 @@ def _get_mounts(
         use_shm_context: bool,
         use_artifacts_context: bool,
         use_tmux_context: bool = False,
+        use_sandbox_context: bool = False,
         run_path: Optional[str] = None,
     ) -> List[VolumeMount]:
         raise NotImplementedError
@@ -681,6 +682,8 @@ def _get_tensorboard_init_container(
     def _get_tools_init_container(
         cls,
         polyaxon_init: "V1PolyaxonInitContainer",
+        use_tmux: bool = True,
+        use_sandbox: bool = False,
     ) -> Container:
         raise NotImplementedError
 
@@ -943,10 +946,14 @@ def get_init_containers(
                 )
             )
 
-        # Add tmux binary
-        if plugins and plugins.tmux:
+        # Add tool binaries
+        if plugins and (plugins.tmux or plugins.sandbox):
             containers.append(
-                self._get_tools_init_container(polyaxon_init=self.polyaxon_init)
+                self._get_tools_init_container(
+                    polyaxon_init=polyaxon_init,
+                    use_tmux=plugins.tmux,
+                    use_sandbox=plugins.sandbox,
+                )
             )
 
         # Add outputs

cli/polyaxon/_sandbox/constants.py

@@ -0,0 +1,2 @@
+SANDBOX_BOOTSTRAP_PATH = "/opt/polyaxon/bin/bootstrap-sandbox.sh"
+SANDBOX_PORT = 9090