Wire sandbox plugin into K8s services

md-mq · md-mq · commit 36f1a1c6d9fd · 2026-05-07T00:55:56.000+02:00
* 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.
diff --git a/cli/polyaxon/_docker/converter/base/mounts.py b/cli/polyaxon/_docker/converter/base/mounts.py
@@ -156,6 +156,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[docker_types.V1VolumeMount]:
         mounts = []
@@ -171,7 +172,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
diff --git a/cli/polyaxon/_env_vars/keys.py b/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"
diff --git a/cli/polyaxon/_flow/plugins/__init__.py b/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
diff --git a/cli/polyaxon/_k8s/converter/base/containers.py b/cli/polyaxon/_k8s/converter/base/containers.py
@@ -35,8 +35,9 @@ def _patch_container(
         container.volume_mounts = to_list(
             container.volume_mounts, check_none=True
         ) + to_list(volume_mounts, check_none=True)
-        container.ports = to_list(container.ports, check_none=True) + to_list(
-            ports, check_none=True
+        container.ports = cls._merge_container_ports(
+            to_list(container.ports, check_none=True),
+            to_list(ports, check_none=True),
         )
         container.resources = container.resources or resources
         container._image_pull_policy = container.image_pull_policy or image_pull_policy
@@ -48,6 +49,41 @@ def _patch_container(
 
         return cls._sanitize_container(container)
 
+    @staticmethod
+    def _normalize_port(port) -> Optional[k8s_schemas.V1ContainerPort]:
+        if isinstance(port, k8s_schemas.V1ContainerPort):
+            return port
+        if isinstance(port, dict):
+            value = port.get("containerPort")
+            if value is None:
+                value = port.get("container_port")
+            if value is None:
+                return None
+            return k8s_schemas.V1ContainerPort(
+                container_port=int(value),
+                name=port.get("name"),
+                host_ip=port.get("hostIP") or port.get("host_ip"),
+                host_port=port.get("hostPort") or port.get("host_port"),
+                protocol=port.get("protocol"),
+            )
+        if isinstance(port, int):
+            return k8s_schemas.V1ContainerPort(container_port=port)
+        return None
+
+    @classmethod
+    def _merge_container_ports(
+        cls, existing: List, incoming: List
+    ) -> List[k8s_schemas.V1ContainerPort]:
+        merged: List[k8s_schemas.V1ContainerPort] = []
+        seen: set = set()
+        for port in list(existing) + list(incoming):
+            normalized = cls._normalize_port(port)
+            if normalized is None or normalized.container_port in seen:
+                continue
+            seen.add(normalized.container_port)
+            merged.append(normalized)
+        return merged
+
     @staticmethod
     def _sanitize_container_env(
         env: List[k8s_schemas.V1EnvVar],
diff --git a/cli/polyaxon/_k8s/converter/base/init.py b/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 = False,
+        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)],
             )
diff --git a/cli/polyaxon/_k8s/converter/base/main.py b/cli/polyaxon/_k8s/converter/base/main.py
@@ -3,9 +3,12 @@
 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
 
 
@@ -33,6 +36,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 +77,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,17 +101,25 @@ 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 = [
-            k8s_schemas.V1ContainerPort(container_port=port)
-            for port in to_list(ports, check_none=True)
-        ]
+        ports = list(to_list(ports, check_none=True))
+        if plugins and plugins.sandbox and SANDBOX_PORT not in ports:
+            ports.append(SANDBOX_PORT)
+
+        ports = [k8s_schemas.V1ContainerPort(container_port=port) for port in ports]
 
         return self._patch_container(
             container=main_container,
diff --git a/cli/polyaxon/_k8s/converter/base/mounts.py b/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
diff --git a/cli/polyaxon/_k8s/converter/converters/service.py b/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,
         )
diff --git a/cli/polyaxon/_k8s/converter/pod/volumes.py b/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
diff --git a/cli/polyaxon/_local_process/converter/base/mounts.py b/cli/polyaxon/_local_process/converter/base/mounts.py
@@ -80,6 +80,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:
         return []
diff --git a/cli/polyaxon/_runner/converter/converter.py b/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 = False,
+        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
diff --git a/cli/polyaxon/_sandbox/constants.py b/cli/polyaxon/_sandbox/constants.py
@@ -0,0 +1,2 @@
+SANDBOX_BOOTSTRAP_PATH = "/opt/polyaxon/bin/bootstrap-sandbox.sh"
+SANDBOX_PORT = 9090
diff --git a/cli/tests/test_k8s/test_converters/test_base/test_containers.py b/cli/tests/test_k8s/test_converters/test_base/test_containers.py
@@ -5,6 +5,30 @@
 from polyaxon._utils.test_utils import BaseTestCase
 
 
+@pytest.mark.k8s_mark
+class TestPatchContainer(BaseTestCase):
+    def test_patch_container_normalizes_and_dedupes_ports(self):
+        existing = k8s_schemas.V1ContainerPort(
+            name="sandbox", container_port=9090
+        )
+        http = k8s_schemas.V1ContainerPort(name="http", container_port=8080)
+        container = ContainerMixin._patch_container(
+            container=k8s_schemas.V1Container(name="main", ports=[existing]),
+            ports=[
+                9090,
+                http,
+                {"name": "http-dup", "containerPort": 8080},
+                {"name": "admin", "containerPort": 8081},
+            ],
+        )
+
+        assert container.ports == [
+            existing,
+            http,
+            k8s_schemas.V1ContainerPort(name="admin", container_port=8081),
+        ]
+
+
 @pytest.mark.k8s_mark
 class TestSanitizeContainerEnv(BaseTestCase):
     def test_sanitize_container_env_value(self):
diff --git a/cli/tests/test_k8s/test_converters/test_base/test_init/test_init_tools.py b/cli/tests/test_k8s/test_converters/test_base/test_init/test_init_tools.py
diff --git a/cli/tests/test_k8s/test_converters/test_base/test_main/test_main_container.py b/cli/tests/test_k8s/test_converters/test_base/test_main/test_main_container.py
diff --git a/cli/tests/test_k8s/test_converters/test_base/test_mounts.py b/cli/tests/test_k8s/test_converters/test_base/test_mounts.py
diff --git a/cli/tests/test_k8s/test_converters/test_converters/test_service_converter.py b/cli/tests/test_k8s/test_converters/test_converters/test_service_converter.py
diff --git a/cli/tests/test_k8s/test_converters/test_pods/test_volumes.py b/cli/tests/test_k8s/test_converters/test_pods/test_volumes.py
