refactor: exec command

Signed-off-by: thxCode <thxcode0824@gmail.com>

Author: thxCode

gpustack_runtime/cmds/deployer.py

@@ -37,6 +37,39 @@
 if TYPE_CHECKING:
     from argparse import Namespace, _SubParsersAction
 
+_IGNORE_ENVS = (
+    "PATH",
+    "HOME",
+    "LANG",
+    "PWD",
+    "SHELL",
+    "LOG",
+    "XDG",
+    "XPC",
+    "SSH",
+    "LC",
+    "LS",
+    "_",
+    "USER",
+    "TERM",
+    "LESS",
+    "SHLVL",
+    "DBUS",
+    "OLDPWD",
+    "MOTD",
+    "LD",
+    "LIB",
+    "PS1",
+    "PY",
+    "VIRTUAL_ENV",
+    "CONDA",
+    "PAGE",
+    "ZSH",
+    "COMMAND_MODE",
+    "TMPDIR",
+    "GPUSTACK_",
+)
+
 
 class CreateRunnerWorkloadSubCommand(SubCommand):
     """
@@ -175,32 +208,7 @@ def run(self):
                 value=value,
             )
             for name, value in os.environ.items()
-            if not name.startswith(
-                (
-                    "PATH",
-                    "HOME",
-                    "LANG",
-                    "PWD",
-                    "SHELL",
-                    "LOG",
-                    "XDG",
-                    "XPC",
-                    "SSH",
-                    "LC",
-                    "LS",
-                    "_",
-                    "USER",
-                    "TERM",
-                    "LESS",
-                    "SHLVL",
-                    "DBUS",
-                    "OLDPWD",
-                    "MOTD",
-                    "LD",
-                    "LIB",
-                    "GPUSTACK_",
-                ),
-            )
+            if not name.startswith(_IGNORE_ENVS)
         ]
         if self.backend:
             resources = ContainerResources(
@@ -377,32 +385,7 @@ def run(self):
                 value=value,
             )
             for name, value in os.environ.items()
-            if not name.startswith(
-                (
-                    "PATH",
-                    "HOME",
-                    "LANG",
-                    "PWD",
-                    "SHELL",
-                    "LOG",
-                    "XDG",
-                    "XPC",
-                    "SSH",
-                    "LC",
-                    "LS",
-                    "_",
-                    "USER",
-                    "TERM",
-                    "LESS",
-                    "SHLVL",
-                    "DBUS",
-                    "OLDPWD",
-                    "MOTD",
-                    "LD",
-                    "LIB",
-                    "GPUSTACK_",
-                ),
-            )
+            if not name.startswith(_IGNORE_ENVS)
         ]
         if self.backend:
             resources = ContainerResources(

gpustack_runtime/deployer/__types__.py

@@ -736,7 +736,7 @@ class WorkloadPlan(WorkloadSecurity):
     """
 
     resource_key_runtime_env_mapping: dict[str, str] = field(
-        default_factory=lambda: envs.GPUSTACK_RUNTIME_DEPLOY_MAP_RUNTIME_VISIBLE_DEVICES,
+        default_factory=lambda: envs.GPUSTACK_RUNTIME_DEPLOY_RESOURCE_KEY_MAP_RUNTIME_VISIBLE_DEVICES,
     )
     """
     Mapping from resource names to environment variable names for device allocation,
@@ -746,7 +746,7 @@ class WorkloadPlan(WorkloadSecurity):
     With privileged mode, the container can access all GPUs even if specified.
     """
     resource_key_backend_env_mapping: dict[str, list[str]] = field(
-        default_factory=lambda: envs.GPUSTACK_RUNTIME_DEPLOY_MAP_BACKEND_VISIBLE_DEVICES,
+        default_factory=lambda: envs.GPUSTACK_RUNTIME_DEPLOY_RESOURCE_KEY_MAP_BACKEND_VISIBLE_DEVICES,
     )
     """
     Mapping from resource names to environment variable names for device runtime,
@@ -918,8 +918,10 @@ class WorkloadStatus:
     Attributes:
         name (WorkloadName):
             Name for the workload, it should be unique in the deployer.
-        created_at (str | None):
+        created_at str:
             Creation time of the workload.
+        namespace (WorkloadNamespace | None):
+            Namespace for the workload.
         labels (dict[str, str] | None):
             Labels for the workload.
         executable (list[WorkloadStatusOperation]):
@@ -940,6 +942,10 @@ class WorkloadStatus:
     """
     Creation time of the workload.
     """
+    namespace: WorkloadNamespace | None = None
+    """
+    Namespace for the workload.
+    """
     labels: dict[str, str] | None = field(default_factory=dict)
     """
     Labels for the workload.
@@ -978,11 +984,11 @@ def fileno(self) -> int:
         raise NotImplementedError
 
     @abstractmethod
-    def read(self, size: int = -1) -> bytes | str | None:
+    def read(self, size: int = -1) -> bytes | None:
         raise NotImplementedError
 
     @abstractmethod
-    def write(self, data: bytes | str) -> int:
+    def write(self, data: bytes) -> int:
         raise NotImplementedError
 
     @abstractmethod
@@ -1016,12 +1022,16 @@ def __init__(self):
 
         if backend := detect_backend():
             rk = envs.GPUSTACK_RUNTIME_DETECT_BACKEND_MAP_RESOURCE_KEY.get(backend)
-            re = envs.GPUSTACK_RUNTIME_DEPLOY_MAP_RUNTIME_VISIBLE_DEVICES.get(rk)
-            be = envs.GPUSTACK_RUNTIME_DEPLOY_MAP_BACKEND_VISIBLE_DEVICES.get(rk)
-            if re:
-                self._runtime_visible_devices_env_name = re
-            if be:
-                self._backend_visible_devices_env_names = be
+            ren = envs.GPUSTACK_RUNTIME_DEPLOY_RESOURCE_KEY_MAP_RUNTIME_VISIBLE_DEVICES.get(
+                rk,
+            )
+            ben = envs.GPUSTACK_RUNTIME_DEPLOY_RESOURCE_KEY_MAP_BACKEND_VISIBLE_DEVICES.get(
+                rk,
+            )
+            if ren:
+                self._runtime_visible_devices_env_name = ren
+            if ben:
+                self._backend_visible_devices_env_names = ben
 
     @staticmethod
     @abstractmethod

gpustack_runtime/deployer/docker.py

@@ -1377,20 +1377,16 @@ def closed(self) -> bool:
         return not (self._sock and not self._sock.closed)
 
     def fileno(self) -> int:
-        if self.closed:
-            return -1
         return self._sock.fileno()
 
-    def read(self, size=-1) -> bytes | str | None:
+    def read(self, size=-1) -> bytes | None:
         if self.closed:
             return None
         return self._sock.read(size)
 
-    def write(self, data: bytes | str) -> int:
+    def write(self, data: bytes) -> int:
         if self.closed:
             return 0
-        if isinstance(data, str):
-            data = data.encode("utf-8")
         return self._sock.write(data)
 
     def close(self):

gpustack_runtime/deployer/kuberentes.py

@@ -188,6 +188,11 @@ def __init__(
         **kwargs,
     ):
         created_at = k_pod.metadata.creation_timestamp.strftime("%Y-%m-%dT%H:%M:%S.%fZ")
+        namespace = (
+            k_pod.metadata.namespace
+            if k_pod.metadata.namespace != envs.GPUSTACK_RUNTIME_KUBERNETES_NAMESPACE
+            else None
+        )
         labels = {
             k: v
             for k, v in (k_pod.metadata.labels or {}).items()
@@ -197,6 +202,7 @@ def __init__(
         super().__init__(
             name=name,
             created_at=created_at,
+            namespace=namespace,
             labels=labels,
             **kwargs,
         )
@@ -1559,25 +1565,34 @@ class KubernetesWorkloadExecStream(WorkloadExecStream):
     def __init__(self, ws: kubernetes.stream.ws_client.WSClient):
         super().__init__()
         self._ws = ws
+        self._ws.run_forever(timeout=1)
+        # self._ws.write_stdin(self._ws.newline)
 
     @property
     def closed(self) -> bool:
         return not (self._ws and self._ws.is_open())
 
     def fileno(self) -> int:
-        if self.closed:
-            return -1
         return self._ws.sock.fileno()
 
-    def read(self, *_) -> bytes | str | None:
+    def recv(self, size=-1) -> bytes | None:
+        return self.read(size)
+
+    def send(self, data: bytes) -> int:
+        return self.write(data)
+
+    def read(self, *_) -> bytes | None:
         if self.closed:
             return None
-        raise NotImplementedError
+        self._ws.update(timeout=1)
+        return self._ws.read_all().encode("utf-8", errors="replace")
 
-    def write(self, *_) -> int:
+    def write(self, data: bytes) -> int:
         if self.closed:
             return 0
-        raise NotImplementedError
+        data_len = len(data)
+        self._ws.write_stdin(data.decode("utf-8", errors="replace"))
+        return data_len
 
     def close(self):
         if not self.closed:

gpustack_runtime/envs.py

@@ -40,16 +40,18 @@
     which is used to tell deployer do a device detection and get the corresponding resource key before mapping.
     e.g., "gpustack.ai/devices".
     """
-    GPUSTACK_RUNTIME_DEPLOY_MAP_RUNTIME_VISIBLE_DEVICES: dict[str, str] | None = None
+    GPUSTACK_RUNTIME_DEPLOY_RESOURCE_KEY_MAP_RUNTIME_VISIBLE_DEVICES: (
+        dict[str, str] | None
+    ) = None
     """
-    Manual mapping of container backend visible devices environment variables,
+    Manual mapping of runtime visible devices environment variables,
     which is used to tell the Container Runtime which GPUs to mount into the container,
     e.g., `{"nvidia.com/devices": "NVIDIA_VISIBLE_DEVICES", "amd.com/devices": "AMD_VISIBLE_DEVICES"}`.
     The key is the resource key, and the value is the environment variable name.
     """
-    GPUSTACK_RUNTIME_DEPLOY_MAP_BACKEND_VISIBLE_DEVICES: dict[str, list[str]] | None = (
-        None
-    )
+    GPUSTACK_RUNTIME_DEPLOY_RESOURCE_KEY_MAP_BACKEND_VISIBLE_DEVICES: (
+        dict[str, list[str]] | None
+    ) = None
     """
     Manual mapping of backend visible devices environment variables,
     which is used to tell the Device Runtime (e.g., ROCm, CUDA, OneAPI) which GPUs to use inside the container,
@@ -137,9 +139,9 @@
         "GPUSTACK_RUNTIME_DEPLOY_AUTOMAP_RESOURCE_KEY",
         "gpustack.ai/devices",
     ),
-    "GPUSTACK_RUNTIME_DEPLOY_MAP_RUNTIME_VISIBLE_DEVICES": lambda: to_dict(
+    "GPUSTACK_RUNTIME_DEPLOY_RESOURCE_KEY_MAP_RUNTIME_VISIBLE_DEVICES": lambda: to_dict(
         getenv(
-            "GPUSTACK_RUNTIME_DEPLOY_MAP_RUNTIME_VISIBLE_DEVICES",
+            "GPUSTACK_RUNTIME_DEPLOY_RESOURCE_KEY_MAP_RUNTIME_VISIBLE_DEVICES",
             "nvidia.com/devices=NVIDIA_VISIBLE_DEVICES;"
             "amd.com/devices=AMD_VISIBLE_DEVICES;"
             "huawei.com/devices=ASCEND_VISIBLE_DEVICES;"
@@ -149,9 +151,9 @@
             "cambricon.com/devices=CAMBRICON_VISIBLE_DEVICES;",
         ),
     ),
-    "GPUSTACK_RUNTIME_DEPLOY_MAP_BACKEND_VISIBLE_DEVICES": lambda: to_dict(
+    "GPUSTACK_RUNTIME_DEPLOY_RESOURCE_KEY_MAP_BACKEND_VISIBLE_DEVICES": lambda: to_dict(
         getenv(
-            "GPUSTACK_RUNTIME_DEPLOY_MAP_BACKEND_VISIBLE_DEVICES",
+            "GPUSTACK_RUNTIME_DEPLOY_RESOURCE_KEY_MAP_BACKEND_VISIBLE_DEVICES",
             "nvidia.com/devices=CUDA_VISIBLE_DEVICES;"
             "amd.com/devices=ROCR_VISIBLE_DEVICES;"
             "huawei.com/devices=ASCEND_RT_VISIBLE_DEVICES,NPU_VISIBLE_DEVICES;"