Add new tmux plugin

Author: polyaxon-ci

cli/polyaxon/_auxiliaries/__init__.py

@@ -17,4 +17,9 @@
     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

@@ -0,0 +1,57 @@
+from typing import Any, Dict, Optional, Union
+
+from clipped.compact.pydantic import Field, StrictStr
+from clipped.types.ref_or_obj import RefField
+from clipped.utils.versions import clean_version_post_suffix
+
+from polyaxon import pkg
+from polyaxon._containers.pull_policy import PullPolicy
+from polyaxon._k8s import k8s_schemas
+from polyaxon._schemas.base import BaseSchemaModel
+
+
+def get_tmux_resources() -> k8s_schemas.V1ResourceRequirements:
+    return k8s_schemas.V1ResourceRequirements(
+        limits={"cpu": "100m", "memory": "64Mi"},
+        requests={"cpu": "50m", "memory": "32Mi"},
+    )
+
+
+class V1PolyaxonTmuxContainer(BaseSchemaModel):
+    _IDENTIFIER = "tmux"
+
+    image: Optional[StrictStr] = None
+    image_tag: Optional[StrictStr] = Field(alias="imageTag", default=None)
+    image_pull_policy: Optional[PullPolicy] = Field(
+        alias="imagePullPolicy", default=None
+    )
+    resources: Optional[Union[Dict[str, Any], RefField]] = None
+
+    def get_image(self):
+        image = self.image or "polyaxon/polyaxon-tmux"
+        image_tag = (
+            self.image_tag
+            if self.image_tag is not None
+            else clean_version_post_suffix(pkg.VERSION)
+        )
+        return "{}:{}".format(image, image_tag) if image_tag else image
+
+    def get_resources(self):
+        return self.resources if self.resources else get_tmux_resources()
+
+
+def get_default_tmux_container(
+    schema=True,
+) -> Union[Dict, V1PolyaxonTmuxContainer]:
+    default = {
+        "image": "polyaxon/polyaxon-tmux",
+        "imageTag": clean_version_post_suffix(pkg.VERSION),
+        "imagePullPolicy": PullPolicy.IF_NOT_PRESENT.value,
+        "resources": {
+            "limits": {"cpu": "100m", "memory": "64Mi"},
+            "requests": {"cpu": "50m", "memory": "32Mi"},
+        },
+    }
+    if schema:
+        return V1PolyaxonTmuxContainer.from_dict(default)
+    return default

cli/polyaxon/_cli/operations.py

@@ -1547,7 +1547,7 @@ def inspect(ctx, project, uid):
     "--command",
     "-cmd",
     type=str,
-    help="Path to download, if not provided the full run's  artifacts will downloaded.",
+    help="The command to execute, default: /bin/bash (or tmux if the plugin is enabled).",
 )
 @click.option(
     "--pod",
@@ -1600,7 +1600,7 @@ def shell(ctx, project, uid, command, pod, container):
         pty = PseudoTerminal(
             client_shell=client.shell(command=command, pod=pod, container=container)
         )
-        pty.start(sys.argv[1:])
+        pty.start()
     except (ApiException, HTTPError) as e:
         handle_cli_error(e, message="Could not inspect the run `{}`.".format(run_uuid))
         sys.exit(1)

cli/polyaxon/_client/run.py

@@ -38,7 +38,7 @@
 from polyaxon._client.decorators import client_handler, get_global_or_inline_config
 from polyaxon._client.mixin import ClientMixin
 from polyaxon._client.store import PolyaxonStore
-from polyaxon._constants.metadata import META_COPY_ARTIFACTS, META_RECOMPILE
+from polyaxon._constants.metadata import META_COPY_ARTIFACTS, META_RECOMPILE, META_TMUX
 from polyaxon._containers.names import MAIN_CONTAINER_NAMES
 from polyaxon._contexts import paths as ctx_paths
 from polyaxon._env_vars.getters import (
@@ -976,6 +976,7 @@ def shell(
                     "The shell command is only usable for operations managed by Polyaxon "
                     "and actively running."
                 )
+            inspection = inspection.get("pods", {})
             if not pod:
                 pod = next(iter(inspection.keys()))
             pod_content = inspection.get(pod, {})
@@ -1027,7 +1028,12 @@ def shell(
         )
         url = absolute_uri(url=url, host=self.client.config.host)
 
-        command = command or "/bin/bash"
+        if not command:
+            command = (
+                "/opt/polyaxon/bin/tmux new-session -A -s terminal"
+                if self._has_meta_key(META_TMUX)
+                else "/bin/bash"
+            )
         return ws_client.websocket_call(
             self.client.config.sdk_config,
             url,

cli/polyaxon/_compiler/resolver/agent.py

@@ -8,8 +8,10 @@
 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
@@ -21,6 +23,7 @@
 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
@@ -50,6 +53,7 @@ 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/_constants/metadata.py

@@ -22,4 +22,5 @@
 META_DESTINATION_IMAGE = "destination_image"
 META_PROGRESS = "progress"
 META_RECOMPILE = "recompile"
+META_TMUX = "tmux"
 META_TTL = "ttl"

cli/polyaxon/_containers/names.py

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

cli/polyaxon/_contexts/paths.py

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

cli/polyaxon/_deploy/schemas/deployment.py

@@ -18,6 +18,7 @@
     V1PolyaxonInitContainer,
     V1PolyaxonNotifier,
     V1PolyaxonSidecarContainer,
+    V1PolyaxonTmuxContainer,
 )
 from polyaxon._auxiliaries.default_scheduling import V1DefaultScheduling
 from polyaxon._connections import V1Connection
@@ -252,6 +253,7 @@ 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/_docker/converter/base/main.py

@@ -58,6 +58,7 @@ def _get_main_container(
                 use_artifacts_context=False,  # Main container has a check and handling for this
                 use_docker_context=plugins.docker,
                 use_shm_context=plugins.shm,
+                use_tmux_context=plugins.tmux,
                 run_path=run_path,
             )
             if plugins