Use single init container for handling all new plugins: tmux, ssh, and sandbox

Author: polyaxon-ci

cli/polyaxon/_auxiliaries/__init__.py

@@ -17,9 +17,4 @@
     get_default_sidecar_container,
     get_sidecar_resources,
 )
-from polyaxon._auxiliaries.tmux import (
-    V1PolyaxonTmuxContainer,
-    get_default_tmux_container,
-    get_tmux_resources,
-)
 from polyaxon._auxiliaries.tuner import get_default_tuner_container

cli/polyaxon/_auxiliaries/tmux.py

@@ -8,10 +8,8 @@
 from polyaxon._auxiliaries import (
     V1PolyaxonInitContainer,
     V1PolyaxonSidecarContainer,
-    V1PolyaxonTmuxContainer,
     get_default_init_container,
     get_default_sidecar_container,
-    get_default_tmux_container,
 )
 from polyaxon._connections import V1Connection, V1ConnectionResource
 from polyaxon._flow import V1CompiledOperation, V1Init
@@ -23,7 +21,6 @@
 class AgentResolver(BaseSchemaModel):
     polyaxon_sidecar: Optional[V1PolyaxonSidecarContainer] = None
     polyaxon_init: Optional[V1PolyaxonInitContainer] = None
-    polyaxon_tmux: Optional[V1PolyaxonTmuxContainer] = None
     namespace: Optional[str] = None
     secrets: Optional[List[Union[V1ConnectionResource, ConnectionResource]]] = None
     config_maps: Optional[List[Union[V1ConnectionResource, ConnectionResource]]] = None
@@ -53,7 +50,6 @@ def resolve(
 
         self.polyaxon_sidecar = agent_config.sidecar or get_default_sidecar_container()
         self.polyaxon_init = agent_config.init or get_default_init_container()
-        self.polyaxon_tmux = agent_config.tmux or get_default_tmux_container()
         if compiled_operation.namespace:
             namespaces = agent_config.additional_namespaces or []
             namespaces.append(agent_config.namespace)

cli/polyaxon/_compiler/resolver/agent.py

@@ -10,7 +10,7 @@
     PYTORCHJOBS_CONTAINER,
 ]
 INIT_AUTH_CONTAINER = "polyaxon-init-auth"
-INIT_TMUX_CONTAINER = "polyaxon-init-tmux"
+INIT_TOOLS_CONTAINER = "polyaxon-init-tools"
 INIT_DOCKERFILE_CONTAINER_PREFIX = "polyaxon-init-dockerfile"
 INIT_FILE_CONTAINER_PREFIX = "polyaxon-init-file"
 INIT_TENSORBOARD_CONTAINER_PREFIX = "polyaxon-init-tb"

cli/polyaxon/_containers/names.py

@@ -42,7 +42,7 @@
     CONTEXT_MOUNT_ARTIFACTS_FORMAT
 )
 CONTEXT_MOUNT_SHM = "/dev/shm"
-CONTEXT_MOUNT_TMUX_BIN = "/opt/polyaxon/bin"
+CONTEXT_MOUNT_TOOLS_BIN = "/opt/polyaxon/bin"
 CONTEXT_MOUNT_DOCKER = "/var/run/docker.sock"
 
 CONTEXT_TMP_POLYAXON_PATH = get_temp_path(".polyaxon")

cli/polyaxon/_contexts/paths.py

@@ -18,7 +18,6 @@
     V1PolyaxonInitContainer,
     V1PolyaxonNotifier,
     V1PolyaxonSidecarContainer,
-    V1PolyaxonTmuxContainer,
 )
 from polyaxon._auxiliaries.default_scheduling import V1DefaultScheduling
 from polyaxon._connections import V1Connection
@@ -253,7 +252,6 @@ class DeploymentConfig(BaseSchemaModel):
     operator: Optional[OperatorServiceConfig] = None
     init: Optional[V1PolyaxonInitContainer] = None
     sidecar: Optional[V1PolyaxonSidecarContainer] = None
-    tmux: Optional[V1PolyaxonTmuxContainer] = None
     notifier: Optional[V1PolyaxonNotifier] = None
     cleaner: Optional[V1PolyaxonCleaner] = None
     default_scheduling: Optional[V1DefaultScheduling] = Field(

cli/polyaxon/_deploy/schemas/deployment.py

@@ -139,12 +139,12 @@ def _get_shm_context_mount(cls) -> docker_types.V1VolumeMount:
         return cls._get_volume(mount_path=ctx_paths.CONTEXT_MOUNT_SHM, read_only=False)
 
     @classmethod
-    def _get_tmux_bin_context_mount(
+    def _get_tools_bin_context_mount(
         cls,
         read_only: bool = True,
     ) -> docker_types.V1VolumeMount:
         return cls._get_volume(
-            mount_path=ctx_paths.CONTEXT_MOUNT_TMUX_BIN,
+            mount_path=ctx_paths.CONTEXT_MOUNT_TOOLS_BIN,
             read_only=read_only,
         )
 
@@ -172,6 +172,6 @@ def _get_mounts(
         if use_shm_context:
             mounts.append(cls._get_shm_context_mount())
         if use_tmux_context:
-            mounts.append(cls._get_tmux_bin_context_mount(read_only=True))
+            mounts.append(cls._get_tools_bin_context_mount(read_only=True))
 
         return mounts

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

@@ -70,7 +70,6 @@
 # Agent
 ENV_KEYS_AGENT_SIDECAR = "POLYAXON_AGENT_SIDECAR"
 ENV_KEYS_AGENT_INIT = "POLYAXON_AGENT_INIT"
-ENV_KEYS_AGENT_TMUX = "POLYAXON_AGENT_TMUX"
 ENV_KEYS_AGENT_CLEANER = "POLYAXON_AGENT_CLEANER"
 ENV_KEYS_AGENT_NOTIFIER = "POLYAXON_AGENT_NOTIFIER"
 ENV_KEYS_AGENT_DEFAULT_SCHEDULING = "POLYAXON_AGENT_DEFAULT_SCHEDULING"

cli/polyaxon/_env_vars/keys.py

@@ -6,7 +6,6 @@
 from polyaxon._auxiliaries import V1PolyaxonInitContainer
 from polyaxon._connections import V1Connection, V1ConnectionKind
 from polyaxon._constants.globals import DEFAULT
-from polyaxon._auxiliaries.tmux import V1PolyaxonTmuxContainer
 from polyaxon._containers.names import (
     INIT_ARTIFACTS_CONTAINER_PREFIX,
     INIT_AUTH_CONTAINER,
@@ -15,7 +14,7 @@
     INIT_FILE_CONTAINER_PREFIX,
     INIT_GIT_CONTAINER_PREFIX,
     INIT_TENSORBOARD_CONTAINER_PREFIX,
-    INIT_TMUX_CONTAINER,
+    INIT_TOOLS_CONTAINER,
     generate_container_name,
 )
 from polyaxon._contexts import paths as ctx_paths
@@ -513,18 +512,18 @@ def _get_auth_context_init_container(
         return cls._patch_container(container)
 
     @classmethod
-    def _get_tmux_init_container(
+    def _get_tools_init_container(
         cls,
-        polyaxon_tmux: V1PolyaxonTmuxContainer,
+        polyaxon_init: V1PolyaxonInitContainer,
     ) -> k8s_schemas.V1Container:
         return cls._patch_container(
             container=k8s_schemas.V1Container(
-                name=INIT_TMUX_CONTAINER,
-                image=polyaxon_tmux.get_image(),
-                image_pull_policy=polyaxon_tmux.image_pull_policy,
+                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"],
-                resources=polyaxon_tmux.get_resources(),
-                volume_mounts=[cls._get_tmux_bin_context_mount(read_only=False)],
+                resources=polyaxon_init.get_resources(),
+                volume_mounts=[cls._get_tools_bin_context_mount(read_only=False)],
             )
         )
 

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

@@ -81,12 +81,12 @@ def _get_shm_context_mount() -> k8s_schemas.V1VolumeMount:
         )
 
     @staticmethod
-    def _get_tmux_bin_context_mount(
+    def _get_tools_bin_context_mount(
         read_only: bool = True,
     ) -> k8s_schemas.V1VolumeMount:
         return k8s_schemas.V1VolumeMount(
-            name=constants.VOLUME_MOUNT_TMUX_BIN,
-            mount_path=ctx_paths.CONTEXT_MOUNT_TMUX_BIN,
+            name=constants.VOLUME_MOUNT_TOOLS_BIN,
+            mount_path=ctx_paths.CONTEXT_MOUNT_TOOLS_BIN,
             read_only=read_only,
         )
 
@@ -114,6 +114,6 @@ def _get_mounts(
         if use_shm_context:
             mounts.append(cls._get_shm_context_mount())
         if use_tmux_context:
-            mounts.append(cls._get_tmux_bin_context_mount(read_only=True))
+            mounts.append(cls._get_tools_bin_context_mount(read_only=True))
 
         return mounts