From aca8d5e990898d64178481375bf33b11cfdcfb29 Mon Sep 17 00:00:00 2001 From: Leonid Pershin Date: Fri, 21 Aug 2026 07:20:28 +0300 Subject: [PATCH] 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. --- src/gpu_rent/remote/civitai_fetch.py | 16 ++-- src/gpu_rent/remote/llamacpp_fetch.py | 13 +-- src/gpu_rent/remote/ollama_pull.py | 4 +- src/gpu_rent/session.py | 3 + src/gpu_rent/ssh_ops.py | 117 +++++++++++++++++++++++--- src/gpu_rent/term.py | 39 ++++++++- tests/test_ssh_stream.py | 33 ++++++++ 7 files changed, 196 insertions(+), 29 deletions(-) create mode 100644 tests/test_ssh_stream.py diff --git a/src/gpu_rent/remote/civitai_fetch.py b/src/gpu_rent/remote/civitai_fetch.py index e221243..ef4c040 100644 --- a/src/gpu_rent/remote/civitai_fetch.py +++ b/src/gpu_rent/remote/civitai_fetch.py @@ -47,7 +47,7 @@ def progress_line( class DownloadProgress: - """Print size + speed about once per second (SSH readline-friendly).""" + """Print size + speed ~1Hz; \\r in place, newline on finish.""" def __init__(self, label: str, total: int | None) -> None: self.label = label @@ -67,14 +67,16 @@ class DownloadProgress: self._emit() def finish(self) -> None: - self._emit() + # Final line with newline so the next log stays below the bar. + self._emit(final=True) - def _emit(self) -> None: + def _emit(self, *, final: bool = False) -> None: elapsed = max(time.monotonic() - self.t0, 0.001) - print( - progress_line(self.label, self.done, self.total, self.done / elapsed), - flush=True, - ) + line = progress_line(self.label, self.done, self.total, self.done / elapsed) + if final: + print(line, flush=True) + else: + print(line, end="\r", flush=True) def sha256_path(path: Path) -> str: diff --git a/src/gpu_rent/remote/llamacpp_fetch.py b/src/gpu_rent/remote/llamacpp_fetch.py index d2dc4ac..e47b735 100644 --- a/src/gpu_rent/remote/llamacpp_fetch.py +++ b/src/gpu_rent/remote/llamacpp_fetch.py @@ -64,14 +64,15 @@ class DownloadProgress: self._emit() def finish(self) -> None: - self._emit() + self._emit(final=True) - def _emit(self) -> None: + def _emit(self, *, final: bool = False) -> None: elapsed = max(time.monotonic() - self.t0, 0.001) - print( - progress_line(self.label, self.done, self.total, self.done / elapsed), - flush=True, - ) + line = progress_line(self.label, self.done, self.total, self.done / elapsed) + if final: + print(line, flush=True) + else: + print(line, end="\r", flush=True) def download(url: str, dest: Path, headers: dict[str, str], *, label: str) -> None: diff --git a/src/gpu_rent/remote/ollama_pull.py b/src/gpu_rent/remote/ollama_pull.py index 7624490..f9d2fd3 100644 --- a/src/gpu_rent/remote/ollama_pull.py +++ b/src/gpu_rent/remote/ollama_pull.py @@ -96,7 +96,9 @@ def pull_stream(name: str, label: str) -> None: if now - last_print >= 1.0 or done >= total: elapsed = max(now - t0, 0.001) speed = done / elapsed - print(progress_line(label, done, total, speed), flush=True) + line = progress_line(label, done, total, speed) + end = "\n" if done >= total else "\r" + print(line, end=end, flush=True) last_print = now elif status and status not in {"success"} and now - last_print >= 2.0: print(f"{label} {status}", flush=True) diff --git a/src/gpu_rent/session.py b/src/gpu_rent/session.py index c68371e..50e5679 100644 --- a/src/gpu_rent/session.py +++ b/src/gpu_rent/session.py @@ -54,6 +54,9 @@ Log = Callable[[str], None] def _log_default(msg: str) -> None: + if msg.startswith("\r"): + print("\r" + msg[1:], end="", flush=True) + return print(msg) diff --git a/src/gpu_rent/ssh_ops.py b/src/gpu_rent/ssh_ops.py index 41b2199..dea0e0f 100644 --- a/src/gpu_rent/ssh_ops.py +++ b/src/gpu_rent/ssh_ops.py @@ -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) diff --git a/src/gpu_rent/term.py b/src/gpu_rent/term.py index f8ccaaa..845587a 100644 --- a/src/gpu_rent/term.py +++ b/src/gpu_rent/term.py @@ -47,6 +47,9 @@ _OK = re.compile( r")" ) +_progress_active = False +_progress_width = 0 + def has_markup(text: str) -> bool: return bool(_MARKUP.search(text)) @@ -80,26 +83,60 @@ def paint(msg: str) -> str: return body +def _end_progress_line() -> None: + global _progress_active, _progress_width + if not _progress_active: + return + sys.stdout.write("\n") + sys.stdout.flush() + _progress_active = False + _progress_width = 0 + + +def _emit_progress(msg: str) -> None: + """Overwrite the current terminal line (download bars).""" + global _progress_active, _progress_width + text = msg.replace("\r", "").replace("\n", "") + pad = max(0, _progress_width - len(text)) + sys.stdout.write("\r" + text + (" " * pad)) + sys.stdout.flush() + _progress_active = True + _progress_width = max(_progress_width, len(text)) + + def log(msg: str = "") -> None: - """CLI log callback — colored print.""" + """CLI log callback — colored print. + + Messages starting with ``\\r`` update one progress line in place + (used by SSH stream for download bars). + """ + if msg.startswith("\r"): + _emit_progress(msg[1:]) + return + _end_progress_line() console.print(paint(msg)) def ok(msg: str) -> None: + _end_progress_line() console.print(f"[green]{escape(msg)}[/green]") def warn(msg: str) -> None: + _end_progress_line() console.print(f"[yellow]{escape(msg)}[/yellow]") def err(msg: str) -> None: + _end_progress_line() console.print(f"[red]{escape(msg)}[/red]") def info(msg: str) -> None: + _end_progress_line() console.print(f"[cyan]{escape(msg)}[/cyan]") def dim(msg: str) -> None: + _end_progress_line() console.print(f"[dim]{escape(msg)}[/dim]") diff --git a/tests/test_ssh_stream.py b/tests/test_ssh_stream.py new file mode 100644 index 0000000..2e0bb4f --- /dev/null +++ b/tests/test_ssh_stream.py @@ -0,0 +1,33 @@ +"""SSH stream splitting for in-place progress bars.""" + +from gpu_rent.ssh_ops import feed_ssh_log, split_ssh_stream + + +def test_split_crlf_line(): + events, rest = split_ssh_stream("hello\r\n") + assert events == [("line", "hello")] + assert rest == "" + + +def test_split_progress_cr(): + events, rest = split_ssh_stream("prog 10%\rprog 20%\r") + assert events == [("progress", "prog 10%"), ("progress", "prog 20%")] + assert rest == "" + + +def test_split_mixed_progress_then_line(): + events, rest = split_ssh_stream("a\rb\r\nok\npartial") + assert events == [ + ("progress", "a"), + ("line", "b"), + ("line", "ok"), + ] + assert rest == "partial" + + +def test_feed_progress_prefix(): + seen = [] + feed_ssh_log(seen.append, "progress", "50%") + assert seen == ["\r50%"] + feed_ssh_log(seen.append, "line", "done") + assert seen[-1] == "done"