feat: support exec

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

Author: thxCode

gpustack_runtime/__main__.py

@@ -9,6 +9,7 @@
     CreateWorkloadSubCommand,
     DeleteWorkloadSubCommand,
     DetectDevicesSubCommand,
+    ExecWorkloadSubCommand,
     GetWorkloadSubCommand,
     ListWorkloadsSubCommand,
     LogsWorkloadSubCommand,
@@ -38,6 +39,7 @@ def main():
     GetWorkloadSubCommand.register(subcommand_parser)
     ListWorkloadsSubCommand.register(subcommand_parser)
     LogsWorkloadSubCommand.register(subcommand_parser)
+    ExecWorkloadSubCommand.register(subcommand_parser)
     DetectDevicesSubCommand.register(subcommand_parser)
 
     # Parse

gpustack_runtime/cmds/__init__.py

@@ -4,6 +4,7 @@
     CreateRunnerWorkloadSubCommand,
     CreateWorkloadSubCommand,
     DeleteWorkloadSubCommand,
+    ExecWorkloadSubCommand,
     GetWorkloadSubCommand,
     ListWorkloadsSubCommand,
     LogsWorkloadSubCommand,
@@ -15,6 +16,7 @@
     "CreateWorkloadSubCommand",
     "DeleteWorkloadSubCommand",
     "DetectDevicesSubCommand",
+    "ExecWorkloadSubCommand",
     "GetWorkloadSubCommand",
     "ListWorkloadsSubCommand",
     "LogsWorkloadSubCommand",

gpustack_runtime/cmds/deployer.py

@@ -4,6 +4,7 @@
 import json
 import os
 import platform
+import sys
 import time
 from argparse import REMAINDER
 from typing import TYPE_CHECKING
@@ -21,6 +22,7 @@
     WorkloadStatus,
     create_workload,
     delete_workload,
+    exec_workload,
     get_workload,
     list_workloads,
     logs_workload,
@@ -567,22 +569,132 @@ def __init__(self, args: Namespace):
     def run(self):
         try:
             print("\033[2J\033[H", end="")
-            logs_stream = logs_workload(self.name, tail=self.tail, follow=self.follow)
+            logs_stream = logs_workload(
+                self.name,
+                tail=self.tail,
+                follow=self.follow,
+            )
             with contextlib.closing(logs_stream) as logs:
                 for line in logs:
                     print(line.decode("utf-8").rstrip())
         except KeyboardInterrupt:
             print("\033[2J\033[H", end="")
 
 
+class ExecWorkloadSubCommand(SubCommand):
+    """
+    Command to execute a command in a workload deployment.
+    """
+
+    name: str
+    command: list[str]
+    interactive: bool = False
+
+    @staticmethod
+    def register(parser: _SubParsersAction):
+        exec_parser = parser.add_parser(
+            "exec",
+            help="execute a command in a workload deployment",
+        )
+
+        exec_parser.add_argument(
+            "name",
+            type=str,
+            help="name of the workload",
+        )
+
+        exec_parser.add_argument(
+            "command",
+            nargs=REMAINDER,
+            help="command to execute in the workload",
+        )
+
+        exec_parser.add_argument(
+            "--interactive",
+            "-i",
+            action="store_true",
+            help="interactive mode",
+        )
+
+        exec_parser.set_defaults(func=ExecWorkloadSubCommand)
+
+    def __init__(self, args: Namespace):
+        self.name = args.name
+        self.command = args.command
+        self.interactive = args.interactive
+
+        if not self.name:
+            msg = "The name argument is required."
+            raise ValueError(msg)
+
+    def run(self):
+        try:
+            if self.interactive:
+                from dockerpty import io, pty  # noqa: PLC0415
+        except ImportError:
+            print(
+                "dockerpty is required for interactive mode. "
+                "Please install it via 'pip install dockerpty'.",
+            )
+            sys.exit(1)
+
+        try:
+            print("\033[2J\033[H", end="")
+            exec_result = exec_workload(
+                self.name,
+                detach=not self.interactive,
+                command=self.command,
+            )
+
+            # Non-interactive mode: print output and exit with the command's exit code
+
+            if not self.interactive:
+                if isinstance(exec_result.output, bytes):
+                    print(exec_result.output.decode("utf-8").rstrip())
+                else:
+                    print(exec_result.output)
+                sys.exit(exec_result.exit_code)
+
+            # Interactive mode: use dockerpty to attach to the exec session
+
+            class ExecOperation(pty.Operation):
+                def __init__(self, socket):
+                    self.stdin = sys.stdin
+                    self.stdout = sys.stdout
+                    self.socket = io.Stream(socket)
+
+                def israw(self, **_):
+                    return self.stdout.isatty()
+
+                def start(self, **_):
+                    stream = self.sockets()
+                    return [
+                        io.Pump(io.Stream(self.stdin), stream, wait_for_output=False),
+                        io.Pump(stream, io.Stream(self.stdout), propagate_close=False),
+                    ]
+
+                def resize(self, height, width, **_):
+                    pass
+
+                def sockets(self):
+                    return self.socket
+
+            exec_op = ExecOperation(exec_result.output)
+            pty.PseudoTerminal(None, exec_op).start()
+        except KeyboardInterrupt:
+            print("\033[2J\033[H", end="")
+
+
 def format_workloads_json(sts: list[WorkloadStatus]) -> str:
     return json.dumps([st.__dict__ for st in sts], indent=2)
 
 
-def format_workloads_table(sts: list[WorkloadStatus], width: int = 100) -> str:
+def format_workloads_table(sts: list[WorkloadStatus]) -> str:
     if not sts:
         return "No workloads found."
 
+    width = 100
+
     headers = ["Name", "State", "Created At"]
     col_widths = [
         len(str(getattr(st, attr.lower().replace(" ", "_"))))

gpustack_runtime/cmds/detector.py

@@ -69,10 +69,12 @@ def format_devices_json(devs: Devices) -> str:
     return json.dumps([dev.__dict__ for dev in devs], indent=2)
 
 
-def format_devices_table(devs: Devices, width: int = 100) -> str:
+def format_devices_table(devs: Devices) -> str:
     if not devs:
         return "No GPUs detected."
 
+    width = 100
+
     # Header section
     dev = devs[0]
     header_content = f"{dev.manufacturer.upper()} "

gpustack_runtime/deployer/__init__.py

@@ -22,6 +22,7 @@
     ContainerSecurity,
     OperationError,
     UnsupportedError,
+    WorkloadExecResult,
     WorkloadOperationToken,
     WorkloadPlan,
     WorkloadSecurity,
@@ -32,6 +33,8 @@
 from .docker import DockerDeployer, DockerWorkloadPlan, DockerWorkloadStatus
 
 if TYPE_CHECKING:
+    from collections.abc import Generator
+
     from .__types__ import Deployer, WorkloadName
 
 deployers: list[Deployer] = [
@@ -156,15 +159,15 @@ def logs_workload(
     tail: int | None = None,
     since: int | None = None,
     follow: bool = False,
-):
+) -> Generator[bytes, None, None] | bytes:
     """
     Get the logs of a workload.
 
     Args:
         name:
             The name of the workload to get logs.
         token:
-            The operation token for authentication.
+            The token for operation.
         timestamps:
             Whether to include timestamps in the logs.
         tail:
@@ -194,6 +197,49 @@ def logs_workload(
     raise UnsupportedError(msg)
 
 
+def exec_workload(
+    name: WorkloadName,
+    token: WorkloadOperationToken | None = None,
+    detach: bool = True,
+    command: list[str] | None = None,
+    args: list[str] | None = None,
+) -> WorkloadExecResult:
+    """
+    Execute a command in a running workload.
+
+    Args:
+        name:
+            The name of the workload to execute the command in.
+        token:
+            The token for operation.
+        detach:
+            Whether to detach from the command execution.
+        command:
+            The command to execute.
+        args:
+            The arguments to pass to the command.
+
+    Returns:
+        If detach is False, return a socket object in the output of WorkloadExecResult.
+            otherwise, return the exit code and output of the command in WorkloadExecResult.
+
+    Raises:
+        UnsupportedError:
+            If no deployer supports the given workload.
+        OperationError:
+            If the deployer fails to execute the command in the workload.
+
+    """
+    for dep in deployers:
+        if not dep.is_supported():
+            continue
+
+        return dep.exec(name, token, detach, command, args)
+
+    msg = "No deployer supports"
+    raise UnsupportedError(msg)
+
+
 __all__ = [
     "Container",
     "ContainerCapabilities",
@@ -216,6 +262,7 @@ def logs_workload(
     "DockerWorkloadStatus",
     "OperationError",
     "UnsupportedError",
+    "WorkloadExecResult",
     "WorkloadOperationToken",
     "WorkloadPlan",
     "WorkloadPlan",
@@ -225,6 +272,7 @@ def logs_workload(
     "WorkloadStatusStateEnum",
     "create_workload",
     "delete_workload",
+    "exec_workload",
     "get_workload",
     "list_workloads",
     "logs_workload",

gpustack_runtime/deployer/__types__.py

@@ -3,9 +3,13 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 from enum import Enum
+from typing import TYPE_CHECKING, NamedTuple
 
 from dataclasses_json import dataclass_json
 
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
 
 class UnsupportedError(Exception):
     """
@@ -788,6 +792,16 @@ class WorkloadStatus:
     """
 
 
+class WorkloadExecResult(NamedTuple):
+    """
+    Result of an exec command.
+
+    """
+
+    exit_code: int | None
+    output: str | bytes | object | None
+
+
 class Deployer(ABC):
     """
     Base class for all deployers.
@@ -895,7 +909,7 @@ def logs(
         tail: int | None = None,
         since: int | None = None,
         follow: bool = False,
-    ):
+    ) -> Generator[bytes, None, None] | bytes:
         """
         Get the logs of a workload.
 
@@ -925,3 +939,42 @@ def logs(
 
         """
         raise NotImplementedError
+
+    @abstractmethod
+    def exec(
+        self,
+        name: WorkloadName,
+        token: WorkloadOperationToken | None = None,
+        detach: bool = True,
+        command: list[str] | None = None,
+        args: list[str] | None = None,
+    ) -> WorkloadExecResult:
+        """
+        Execute a command in a workload.
+
+        Args:
+            name:
+                The name of the workload.
+            token:
+                The operation token of the workload.
+                If not specified, execute in the first executable container.
+            detach:
+                Whether to detach from the command.
+            command:
+                The command to execute.
+                If not specified, use the command defined in the container.
+            args:
+                The arguments to pass to the command.
+
+        Returns:
+            If detach is False, return a socket object in the output of WorkloadExecResult.
+            otherwise, return the exit code and output of the command in WorkloadExecResult.
+
+        Raises:
+            UnsupportedError:
+                If the deployer is not supported in the current environment.
+            OperationError:
+                If the workload fails to execute the command.
+
+        """
+        raise NotImplementedError