Skip to content

Commit 06e9ae1

Browse files
V2arKclaude
andcommitted
feat: add CLI shell and exec commands for deployment pod terminal access
Add two new commands under `centml cluster`: - `shell <id>` -- interactive terminal session (like docker exec -it) - `exec <id> -- <command>` -- run a command and return output (like ssh host "cmd") Both connect via WebSocket through the Platform API terminal proxy, matching the same protocol used by the Web UI TerminalView. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 470a8b6 commit 06e9ae1

5 files changed

Lines changed: 745 additions & 0 deletions

File tree

‎centml/cli/main.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
from centml.cli.login import login, logout
44
from centml.cli.cluster import ls, get, delete, pause, resume
5+
from centml.cli.shell import shell, exec_cmd
56

67

78
@click.group()
@@ -47,6 +48,8 @@ def ccluster():
4748
ccluster.add_command(delete)
4849
ccluster.add_command(pause)
4950
ccluster.add_command(resume)
51+
ccluster.add_command(shell)
52+
ccluster.add_command(exec_cmd, name="exec")
5053

5154

5255
cli.add_command(ccluster, name="cluster")

‎centml/cli/shell.py‎

Lines changed: 309 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,309 @@
1+
"""CLI commands for interactive shell and command execution in deployment pods."""
2+
3+
import asyncio
4+
import json
5+
import re
6+
import shutil
7+
import signal
8+
import sys
9+
import urllib.parse
10+
11+
import click
12+
13+
from centml.sdk import PodStatus
14+
from centml.sdk import auth
15+
from centml.sdk.api import get_centml_client
16+
from centml.sdk.config import settings
17+
from centml.cli.cluster import handle_exception
18+
19+
# Lazy-import to keep module loadable without websockets installed at import time,
20+
# and to allow tests to patch the module attribute easily.
21+
import websockets
22+
23+
# These are only available on Unix; guarded at command level via isatty check.
24+
import termios
25+
import tty
26+
27+
28+
def _build_ws_url(api_url, deployment_id, pod_name, shell=None):
29+
"""Build the WebSocket URL for a terminal connection."""
30+
ws_base = api_url.replace("https://", "wss://").replace("http://", "ws://")
31+
url = f"{ws_base}/deployments/{deployment_id}/terminal?pod={urllib.parse.quote(pod_name)}"
32+
if shell:
33+
url += f"&shell={urllib.parse.quote(shell)}"
34+
return url
35+
36+
37+
def _resolve_pod(cclient, deployment_id, pod_name=None):
38+
"""Resolve which pod to connect to.
39+
40+
Args:
41+
cclient: CentMLClient instance.
42+
deployment_id: The deployment ID.
43+
pod_name: Optional specific pod name to target.
44+
45+
Returns:
46+
The pod name string to connect to.
47+
48+
Raises:
49+
click.ClickException: If no running pods or specified pod not found.
50+
"""
51+
status = cclient.get_status_v3(deployment_id)
52+
running_pods = []
53+
for revision in (status.revision_pod_details_list or []):
54+
for pod in (revision.pod_details_list or []):
55+
if pod.status == PodStatus.RUNNING and pod.name:
56+
running_pods.append(pod.name)
57+
58+
if not running_pods:
59+
raise click.ClickException(
60+
f"No running pods found for deployment {deployment_id}"
61+
)
62+
63+
if pod_name is not None:
64+
if pod_name not in running_pods:
65+
pods_list = ", ".join(running_pods)
66+
raise click.ClickException(
67+
f"Pod '{pod_name}' not found. Available running pods: {pods_list}"
68+
)
69+
return pod_name
70+
71+
if len(running_pods) > 1:
72+
click.echo(
73+
f"Multiple running pods found, connecting to {running_pods[0]}. "
74+
f"Use --pod to specify a different pod.",
75+
err=True,
76+
)
77+
return running_pods[0]
78+
79+
80+
async def _forward_io(ws):
81+
"""Bidirectional forwarding between local stdin/stdout and WebSocket.
82+
83+
Returns the remote exit code.
84+
"""
85+
loop = asyncio.get_running_loop()
86+
exit_code = 0
87+
stdin_fd = sys.stdin.fileno()
88+
89+
stdin_closed = asyncio.Event()
90+
91+
async def _read_ws():
92+
nonlocal exit_code
93+
async for raw_msg in ws:
94+
msg = json.loads(raw_msg)
95+
if msg.get("data"):
96+
sys.stdout.buffer.write(msg["data"].encode("utf-8", errors="replace"))
97+
sys.stdout.buffer.flush()
98+
elif msg.get("error"):
99+
sys.stderr.write(f"Error: {msg['error']}\n")
100+
sys.stderr.flush()
101+
if "Code" in msg:
102+
exit_code = msg["Code"]
103+
return
104+
105+
async def _read_stdin():
106+
read_queue = asyncio.Queue()
107+
108+
def _on_stdin_ready():
109+
data = sys.stdin.buffer.read1(4096)
110+
if data:
111+
read_queue.put_nowait(data)
112+
else:
113+
stdin_closed.set()
114+
115+
loop.add_reader(stdin_fd, _on_stdin_ready)
116+
try:
117+
while not stdin_closed.is_set():
118+
try:
119+
data = await asyncio.wait_for(read_queue.get(), timeout=0.5)
120+
except asyncio.TimeoutError:
121+
continue
122+
rows, cols = shutil.get_terminal_size()
123+
await ws.send(json.dumps({
124+
"operation": "stdin",
125+
"data": data.decode("utf-8", errors="replace"),
126+
"rows": rows,
127+
"cols": cols,
128+
}))
129+
finally:
130+
loop.remove_reader(stdin_fd)
131+
132+
tasks = [
133+
asyncio.create_task(_read_ws()),
134+
asyncio.create_task(_read_stdin()),
135+
]
136+
done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
137+
for t in pending:
138+
t.cancel()
139+
for t in done:
140+
if t.exception() is not None:
141+
raise t.exception()
142+
return exit_code
143+
144+
145+
async def _interactive_session(ws_url, token):
146+
"""Run an interactive terminal session over WebSocket.
147+
148+
Enters raw mode, forwards I/O bidirectionally, and restores terminal on exit.
149+
"""
150+
fd = sys.stdin.fileno()
151+
old_settings = termios.tcgetattr(fd)
152+
try:
153+
tty.setraw(fd)
154+
rows, cols = shutil.get_terminal_size()
155+
156+
headers = {"Authorization": f"Bearer {token}"}
157+
async with websockets.connect(ws_url, additional_headers=headers) as ws:
158+
await ws.send(json.dumps({
159+
"operation": "resize",
160+
"rows": rows,
161+
"cols": cols,
162+
}))
163+
164+
loop = asyncio.get_running_loop()
165+
166+
def _on_resize():
167+
r, c = shutil.get_terminal_size()
168+
asyncio.ensure_future(ws.send(json.dumps({
169+
"operation": "resize",
170+
"rows": r,
171+
"cols": c,
172+
})))
173+
174+
loop.add_signal_handler(signal.SIGWINCH, _on_resize)
175+
176+
try:
177+
exit_code = await _forward_io(ws)
178+
finally:
179+
loop.remove_signal_handler(signal.SIGWINCH)
180+
181+
return exit_code
182+
finally:
183+
termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
184+
185+
186+
_ANSI_ESCAPE_RE = re.compile(r"\x1b\[[0-9;?]*[a-zA-Z]|\x1b\].*?\x07|\x1b\[.*?\x1b\\")
187+
188+
_BEGIN_MARKER = "__CENTML_BEGIN_5f3a__"
189+
_END_MARKER = "__CENTML_END_5f3a__"
190+
191+
# printf octal \137 = underscore. The decoded output matches _BEGIN/_END_MARKER,
192+
# but the literal command text does NOT, so shell echo won't trigger false matches.
193+
_PRINTF_BEGIN = r"\137\137CENTML_BEGIN_5f3a\137\137"
194+
_PRINTF_END = r"\137\137CENTML_END_5f3a\137\137"
195+
196+
197+
def _strip_ansi(text):
198+
"""Remove ANSI escape sequences from text."""
199+
return _ANSI_ESCAPE_RE.sub("", text)
200+
201+
202+
async def _exec_session(ws_url, token, command):
203+
"""Execute a command in a pod and return its exit code.
204+
205+
Does not enter raw mode -- output is pipe-friendly.
206+
Suppresses shell echo and uses markers to capture only command output.
207+
"""
208+
rows, cols = shutil.get_terminal_size(fallback=(80, 24))
209+
headers = {"Authorization": f"Bearer {token}"}
210+
211+
async with websockets.connect(ws_url, additional_headers=headers) as ws:
212+
await ws.send(json.dumps({
213+
"operation": "resize",
214+
"rows": rows,
215+
"cols": cols,
216+
}))
217+
218+
# Suppress echo/bracketed-paste, emit begin marker, run command,
219+
# emit end marker with exit code, then exit.
220+
# Markers use printf octal escapes so the literal marker string
221+
# doesn't appear in the command echo.
222+
wrapped = (
223+
f"stty -echo 2>/dev/null; printf '\\033[?2004l';"
224+
f" printf '{_PRINTF_BEGIN}\\n';"
225+
f" {command};"
226+
f" __ec=$?;"
227+
f" printf '\\n{_PRINTF_END}:%d\\n' \"$__ec\";"
228+
f" exit $__ec\n"
229+
)
230+
231+
await ws.send(json.dumps({
232+
"operation": "stdin",
233+
"data": wrapped,
234+
"rows": rows,
235+
"cols": cols,
236+
}))
237+
238+
exit_code = 0
239+
buffer = ""
240+
is_capturing = False
241+
async for raw_msg in ws:
242+
msg = json.loads(raw_msg)
243+
if msg.get("data"):
244+
buffer += msg["data"]
245+
# Process complete lines from buffer
246+
while "\n" in buffer:
247+
line, buffer = buffer.split("\n", 1)
248+
clean = _strip_ansi(line).rstrip("\r")
249+
if _BEGIN_MARKER in clean:
250+
is_capturing = True
251+
continue
252+
if _END_MARKER in clean:
253+
# Parse exit code from marker line
254+
parts = clean.split(_END_MARKER + ":")
255+
if len(parts) > 1:
256+
try:
257+
exit_code = int(parts[1].strip())
258+
except ValueError:
259+
pass
260+
is_capturing = False
261+
continue
262+
if is_capturing:
263+
sys.stdout.write(line + "\n")
264+
sys.stdout.flush()
265+
elif msg.get("error"):
266+
sys.stderr.write(f"Error: {msg['error']}\n")
267+
return 1
268+
if "Code" in msg:
269+
exit_code = msg["Code"]
270+
break
271+
return exit_code
272+
273+
274+
@click.command(help="Open an interactive shell to a deployment pod")
275+
@click.argument("deployment_id", type=int)
276+
@click.option("--pod", default=None, help="Specific pod name (auto-selects first running pod)")
277+
@click.option("--shell", "shell_type", default=None,
278+
type=click.Choice(["bash", "sh", "zsh"]), help="Shell type")
279+
@handle_exception
280+
def shell(deployment_id, pod, shell_type):
281+
if not sys.stdin.isatty():
282+
raise click.ClickException("Interactive shell requires a terminal (TTY)")
283+
284+
with get_centml_client() as cclient:
285+
pod_name = _resolve_pod(cclient, deployment_id, pod)
286+
287+
ws_url = _build_ws_url(settings.CENTML_PLATFORM_API_URL, deployment_id, pod_name, shell_type)
288+
token = auth.get_centml_token()
289+
exit_code = asyncio.run(_interactive_session(ws_url, token))
290+
sys.exit(exit_code)
291+
292+
293+
@click.command(help="Execute a command in a deployment pod",
294+
context_settings=dict(ignore_unknown_options=True))
295+
@click.argument("deployment_id", type=int)
296+
@click.argument("command", nargs=-1, required=True, type=click.UNPROCESSED)
297+
@click.option("--pod", default=None, help="Specific pod name")
298+
@click.option("--shell", "shell_type", default=None,
299+
type=click.Choice(["bash", "sh", "zsh"]), help="Shell type")
300+
@handle_exception
301+
def exec_cmd(deployment_id, command, pod, shell_type):
302+
with get_centml_client() as cclient:
303+
pod_name = _resolve_pod(cclient, deployment_id, pod)
304+
305+
ws_url = _build_ws_url(settings.CENTML_PLATFORM_API_URL, deployment_id, pod_name, shell_type)
306+
token = auth.get_centml_token()
307+
cmd_str = " ".join(command)
308+
exit_code = asyncio.run(_exec_session(ws_url, token, cmd_str))
309+
sys.exit(exit_code)

‎centml/sdk/api.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,9 @@ def get(self, depl_type):
2727
def get_status(self, id):
2828
return self._api.get_deployment_status_deployments_status_deployment_id_get(id)
2929

30+
def get_status_v3(self, deployment_id):
31+
return self._api.get_deployment_status_v3_deployments_status_v3_deployment_id_get(deployment_id)
32+
3033
def get_inference(self, id):
3134
"""Get Inference deployment details - automatically handles both V2 and V3 deployments"""
3235
# Try V3 first (recommended), fallback to V2 if deployment is V2

‎requirements.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,4 +9,5 @@ cryptography==44.0.1
99
prometheus-client>=0.20.0
1010
scipy>=1.6.0
1111
scikit-learn>=1.5.1
12+
websockets>=13.0
1213
platform-api-python-client==4.6.0

0 commit comments

Comments
 (0)