Enhance logging and progress handling in SSH operations
- Updated the `_log_default` function to handle messages starting with `\r` for in-place updates. - Introduced `split_ssh_stream` to differentiate between line and progress events in SSH output. - Implemented `feed_ssh_log` to dispatch log messages based on event type. - Refactored `_stream_pty_output` to utilize the new logging and event handling functions. - Enhanced the `log` function in `term.py` to support progress updates, ensuring a cleaner output experience. - Updated `DownloadProgress` classes in `civitai_fetch.py` and `llamacpp_fetch.py` to print progress in place, improving user feedback during downloads.
This commit is contained in:
+103
-14
@@ -15,6 +15,107 @@ from gpu_rent.config import Config
|
||||
from gpu_rent.errors import CloudError
|
||||
|
||||
|
||||
def split_ssh_stream(buf: str) -> tuple[list[tuple[str, str]], str]:
|
||||
"""Split SSH/PTY stdout into (kind, text) events.
|
||||
|
||||
kind is ``line`` (ended with \\n or \\r\\n) or ``progress`` (ended with lone \\r).
|
||||
Returns (events, remainder).
|
||||
"""
|
||||
events: list[tuple[str, str]] = []
|
||||
i = 0
|
||||
start = 0
|
||||
n = len(buf)
|
||||
while i < n:
|
||||
ch = buf[i]
|
||||
if ch == "\r":
|
||||
if i + 1 < n and buf[i + 1] == "\n":
|
||||
events.append(("line", buf[start:i]))
|
||||
i += 2
|
||||
start = i
|
||||
continue
|
||||
events.append(("progress", buf[start:i]))
|
||||
i += 1
|
||||
start = i
|
||||
continue
|
||||
if ch == "\n":
|
||||
events.append(("line", buf[start:i]))
|
||||
i += 1
|
||||
start = i
|
||||
continue
|
||||
i += 1
|
||||
return events, buf[start:]
|
||||
|
||||
|
||||
def feed_ssh_log(
|
||||
log: Callable[[str], None] | None,
|
||||
kind: str,
|
||||
text: str,
|
||||
) -> None:
|
||||
"""Dispatch a stream event to the CLI log callback."""
|
||||
if not log or not text:
|
||||
# Empty progress is a no-op; empty line still logs blank via caller if needed.
|
||||
if log and kind == "line" and text == "":
|
||||
log("")
|
||||
return
|
||||
if kind == "progress":
|
||||
log("\r" + text)
|
||||
else:
|
||||
log(text)
|
||||
|
||||
|
||||
def _stream_pty_output(
|
||||
stdout,
|
||||
*,
|
||||
log: Callable[[str], None] | None,
|
||||
chunks: list[str],
|
||||
) -> None:
|
||||
"""Read PTY stdout in chunks; honor \\r progress and \\n lines."""
|
||||
buf = ""
|
||||
channel = getattr(stdout, "channel", None)
|
||||
if channel is None:
|
||||
while True:
|
||||
piece = stdout.read(4096)
|
||||
if not piece:
|
||||
break
|
||||
text = piece.decode("utf-8", errors="replace") if isinstance(piece, bytes) else piece
|
||||
chunks.append(text)
|
||||
buf += text
|
||||
events, buf = split_ssh_stream(buf)
|
||||
for kind, part in events:
|
||||
feed_ssh_log(log, kind, part)
|
||||
if buf:
|
||||
feed_ssh_log(log, "line", buf.rstrip("\r"))
|
||||
return
|
||||
|
||||
while True:
|
||||
if channel.recv_ready():
|
||||
data = channel.recv(4096)
|
||||
if not data:
|
||||
break
|
||||
text = data.decode("utf-8", errors="replace")
|
||||
chunks.append(text)
|
||||
buf += text
|
||||
events, buf = split_ssh_stream(buf)
|
||||
for kind, part in events:
|
||||
feed_ssh_log(log, kind, part)
|
||||
continue
|
||||
if channel.exit_status_ready():
|
||||
while channel.recv_ready():
|
||||
data = channel.recv(4096)
|
||||
if not data:
|
||||
break
|
||||
text = data.decode("utf-8", errors="replace")
|
||||
chunks.append(text)
|
||||
buf += text
|
||||
events, buf = split_ssh_stream(buf)
|
||||
for kind, part in events:
|
||||
feed_ssh_log(log, kind, part)
|
||||
break
|
||||
time.sleep(0.05)
|
||||
if buf:
|
||||
feed_ssh_log(log, "line", buf.rstrip("\r"))
|
||||
|
||||
|
||||
def wait_tcp(host: str, port: int, timeout: float = 300.0) -> None:
|
||||
deadline = time.time() + timeout
|
||||
last = None
|
||||
@@ -219,13 +320,7 @@ def run_script_sudo(
|
||||
prefix = f"sudo -n env {env_s} " if env_s else "sudo -n "
|
||||
command = f"{prefix}bash {remote_path}"
|
||||
_stdin, stdout, stderr = client.exec_command(command, timeout=timeout, get_pty=True)
|
||||
while True:
|
||||
line = stdout.readline()
|
||||
if not line:
|
||||
break
|
||||
chunks.append(line)
|
||||
if log:
|
||||
log(line.rstrip("\n\r"))
|
||||
_stream_pty_output(stdout, log=log, chunks=chunks)
|
||||
code = stdout.channel.recv_exit_status()
|
||||
err = stderr.read().decode("utf-8", errors="replace") if not stdout.channel.closed else ""
|
||||
out = "".join(chunks)
|
||||
@@ -256,13 +351,7 @@ def run_python(
|
||||
sftp.close()
|
||||
command = f"python3 {shlex.quote(remote_path)}"
|
||||
_stdin, stdout, stderr = client.exec_command(command, timeout=timeout, get_pty=True)
|
||||
while True:
|
||||
line = stdout.readline()
|
||||
if not line:
|
||||
break
|
||||
chunks.append(line)
|
||||
if log:
|
||||
log(line.rstrip("\n\r"))
|
||||
_stream_pty_output(stdout, log=log, chunks=chunks)
|
||||
code = stdout.channel.recv_exit_status()
|
||||
err = stderr.read().decode("utf-8", errors="replace") if not stdout.channel.closed else ""
|
||||
out = "".join(chunks)
|
||||
|
||||
Reference in New Issue
Block a user