llm-trainer/backend/ssh_client.py

import base64
import threading
import time
from typing import Optional

import paramiko


class SSHClient:
    def __init__(self):
        self.client: Optional[paramiko.SSHClient] = None
        self.connected = False
        self.host = ""
        self.username = ""
        self.port = 22
        self._keepalive_thread: Optional[threading.Thread] = None
        self._stop_keepalive = threading.Event()
        self._lock = threading.Lock()

    def connect(self, host: str, username: str, password: str = None,
                key_path: str = None, port: int = 22) -> bool:
        with self._lock:
            try:
                if self.client:
                    self.client.close()

                self.client = paramiko.SSHClient()
                self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

                kwargs = {"hostname": host, "port": port, "username": username, "timeout": 10}
                if key_path:
                    kwargs["key_filename"] = key_path
                if password:
                    kwargs["password"] = password

                self.client.connect(**kwargs)
                self.connected = True
                self.host = host
                self.username = username
                self.port = port

                self._stop_keepalive.clear()
                self._keepalive_thread = threading.Thread(target=self._keepalive_loop, daemon=True)
                self._keepalive_thread.start()
                return True
            except Exception as e:
                self.connected = False
                raise e

    def disconnect(self):
        self._stop_keepalive.set()
        if self.client:
            self.client.close()
        self.connected = False

    def _keepalive_loop(self):
        while not self._stop_keepalive.wait(30):
            try:
                transport = self.client.get_transport()
                if transport and transport.is_active():
                    transport.send_ignore()
                else:
                    self.connected = False
                    break
            except Exception:
                self.connected = False
                break

    def execute(self, command: str, use_conda: bool = True) -> tuple:
        if not self.is_connected():
            raise Exception("Not connected to SSH server")

        if use_conda:
            full_cmd = (
                f"source /home/{self.username}/miniconda3/etc/profile.d/conda.sh && "
                f"conda activate synthetic-data && {command}"
            )
        else:
            full_cmd = command

        _, stdout, stderr = self.client.exec_command(full_cmd)
        out = stdout.read().decode("utf-8", errors="replace")
        err = stderr.read().decode("utf-8", errors="replace")
        exit_code = stdout.channel.recv_exit_status()
        return out, err, exit_code

    def execute_stream(self, command: str, use_conda: bool = True):
        """Generator that yields output lines from a command."""
        if not self.is_connected():
            raise Exception("Not connected to SSH server")

        if use_conda:
            full_cmd = (
                f"source /home/{self.username}/miniconda3/etc/profile.d/conda.sh && "
                f"conda activate synthetic-data && {command}"
            )
        else:
            full_cmd = command

        transport = self.client.get_transport()
        channel = transport.open_session()
        channel.get_pty()
        channel.exec_command(full_cmd)

        buffer = b""
        while True:
            if channel.recv_ready():
                data = channel.recv(4096)
                if not data:
                    break
                buffer += data
                while b"\n" in buffer:
                    line, buffer = buffer.split(b"\n", 1)
                    yield line.decode("utf-8", errors="replace") + "\n"
            elif channel.exit_status_ready():
                if buffer:
                    yield buffer.decode("utf-8", errors="replace")
                break
            else:
                time.sleep(0.05)

        channel.close()

    def open_shell_channel(self, term: str = "xterm-256color", width: int = 220, height: int = 50):
        """Open an interactive shell channel for the terminal panel."""
        if not self.is_connected():
            raise Exception("Not connected to SSH server")

        transport = self.client.get_transport()
        channel = transport.open_session()
        channel.get_pty(term=term, width=width, height=height)
        channel.invoke_shell()

        # Auto-activate conda env
        activate = (
            f"source /home/{self.username}/miniconda3/etc/profile.d/conda.sh && "
            f"conda activate synthetic-data\n"
        )
        channel.send(activate)
        return channel

    def upload_file(self, local_path: str, remote_path: str):
        if not self.is_connected():
            raise Exception("Not connected to SSH server")
        sftp = self.client.open_sftp()
        try:
            sftp.put(local_path, remote_path)
        finally:
            sftp.close()

    def read_remote_file(self, remote_path: str) -> str:
        out, err, code = self.execute(f"cat '{remote_path}'", use_conda=False)
        if code != 0:
            raise Exception(f"Failed to read file: {err}")
        return out

    def write_remote_file(self, remote_path: str, content: str):
        encoded = base64.b64encode(content.encode()).decode()
        cmd = f"echo '{encoded}' | base64 -d > '{remote_path}'"
        out, err, code = self.execute(cmd, use_conda=False)
        if code != 0:
            raise Exception(f"Failed to write file: {err}")

    def is_connected(self) -> bool:
        try:
            if self.client:
                transport = self.client.get_transport()
                if transport and transport.is_active():
                    return True
        except Exception:
            pass
        self.connected = False
        return False


# Singleton shared across all routes
ssh_manager = SSHClient()