Skip to content

SSH Utils

capture_tmux_pane(ssh, tmux_session_name='vllm_server', verbose=True)

Capture the content of the specified tmux pane.

Parameters:

Name Type Description Default
ssh SSHClient

The SSH client for the connection.

 required
tmux_session_name str

The name of the tmux session. Defaults to "vllm_server".

 'vllm_server'
verbose bool

If True, display verbose output. Defaults to True.

 True

Returns:

Name Type Description
tuple

A tuple containing stdin, stdout, and stderr from the SSH command execution.
Source code in app/utils/ssh_utils.py 
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
def capture_tmux_pane(
    ssh: paramiko.SSHClient,
    tmux_session_name: str = "vllm_server",
    verbose: bool =True
):
    """
    Capture the content of the specified tmux pane.

    Args:
        ssh (paramiko.SSHClient): The SSH client for the connection.
        tmux_session_name (str, optional): The name of the tmux session. Defaults to "vllm_server".
        verbose (bool, optional): If True, display verbose output. Defaults to True.

    Returns:
        tuple: A tuple containing stdin, stdout, and stderr from the SSH command execution.
    """
    ssh_stdin, ssh_stdout, ssh_stderr = execute_ssh_command(ssh,
        f"tmux capture-pane -p -t {tmux_session_name}.0", verbose)
    return ssh_stdin, ssh_stdout, ssh_stderr

execute_ssh_command(ssh, cmd, verbose=True)

Runs the given command in paramiko ssh client

Parameters:

Name Type Description Default
ssh SSHClient

The SSH client for the connection.

 required
cmd str

Command that can run in tmux session terminal.

 required
Source code in app/utils/ssh_utils.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def execute_ssh_command(
        ssh: paramiko.SSHClient,
        cmd: str,
        verbose: bool =True
    ):
    """Runs the given command in paramiko ssh client 

    Args:
        ssh (paramiko.SSHClient): The SSH client for the connection.
        cmd (str): Command that can run in tmux session terminal.
    """
    if verbose: print( f"ssh : {cmd}" )
    ssh_stdin, ssh_stdout, ssh_stderr = ssh.exec_command(cmd)
    ssh_stdout_str = ssh_stdout.read().decode()
    ssh_stderr_str = ssh_stderr.read().decode()
    if verbose:
        if ssh_stdout_str:
            print("\t stdout : ", ssh_stdout_str)
        else:
            print("\t stderr : ", ssh_stderr_str)
    return ssh_stdin, ssh_stdout_str, ssh_stderr_str

get_ssh_session(ip_address, timeout=15, key_path='app/keys/connection-key.pem')

Get paramiko SSHClient

Parameters:

Name Type Description Default
ip_address str

ip address

 required
timeout int

Defaults to 15.

 15
key_path str)

Path to the pem file required to connect to ec2 instance.

 'app/keys/connection-key.pem'
Source code in app/utils/ssh_utils.py 
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
def get_ssh_session(
    ip_address: str, 
    timeout:int =15,
    key_path:str ="app/keys/connection-key.pem"
):
    """Get paramiko SSHClient  

    Args:
        ip_address (str): ip address
        timeout (int, optional): Defaults to 15.
        key_path (str) : Path to the pem file required to connect to ec2 instance. 
    """
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    ssh.connect(ip_address, 
                username="ubuntu", 
                key_filename=key_path,
                timeout=timeout)
    return ssh

is_running_vllm_server(ssh, tmux_session_name='vllm_server', verbose=True)

Check if vllm server has started successfully.

Parameters:

Name Type Description Default
ssh SSHClient

(paramiko.SSHClient): The SSH client for the connection.

 required
tmux_session_name str

Defaults to "vllm_server".

 'vllm_server'
Source code in app/utils/ssh_utils.py 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
def is_running_vllm_server(
        ssh: paramiko.SSHClient,
        tmux_session_name: str = "vllm_server",
        verbose: bool =True
    ):
    """Check if vllm server has started successfully.  

    Args:
        ssh: (paramiko.SSHClient):  The SSH client for the connection.
        tmux_session_name (str, optional): Defaults to "vllm_server".
    """
    ssh_stdin, ssh_stdout, ssh_stderr = capture_tmux_pane(ssh, tmux_session_name, verbose)
    if "Uvicorn running on http" in ssh_stdout:
        return True, "running"
    elif "python -m vllm" in ssh_stdout or \
            "llm_engine.py" in ssh_stdout:
        return False, "loading"
    else:
        return False, ""

start_vllm_server(ssh, model_id, ip_address, port=8000, conda_env_name='pytorch', tmux_session_name='vllm_server')

Starts a VLLM Server for the provided Model ID.

Parameters:

Name Type Description Default
ssh SSHClient

The SSH client for the connection.

 required
model_id str

HF supported model ID. TODO: Support s3 checkpoint path downloading.

 required
ip_address str

The IP address for the EC2 instance.

 required
port str

VLLM runs a server that listens on this port. Ensure that this port is open for requests.

 8000
conda_env_name str

The specified environment should have VLLM installed.

 'pytorch'
tmux_session_name str

This is the name of the tmux session that runs the VLLM server.

 'vllm_server'
Source code in app/utils/ssh_utils.py 
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def start_vllm_server(
        ssh: paramiko.SSHClient,
        model_id: str,
        ip_address: str,
        port: int =8000,
        conda_env_name: str = "pytorch",
        tmux_session_name: str = "vllm_server"
    ):

    """Starts a VLLM Server for the provided Model ID.

    Args:
        ssh (paramiko.SSHClient):  The SSH client for the connection.
        model_id (str): HF supported model ID. TODO: Support s3 checkpoint path downloading.
        ip_address (str): The IP address for the EC2 instance.
        port (str): VLLM runs a server that listens on this port. Ensure that this port is open for requests.
        conda_env_name (str): The specified environment should have VLLM installed.
        tmux_session_name (str): This is the name of the tmux session that runs the VLLM server.
    """

    """check if session is already created"""
    _, ssh_stdout, _ = execute_ssh_command(ssh,"tmux ls")
    if tmux_session_name not in ssh_stdout:
        """start a new tmux session"""
        _,_,_ = execute_ssh_command(ssh,  
                    f"tmux new -d -s {tmux_session_name}" )
        time.sleep(5)

        """activate environment"""
        _, _, _ = execute_ssh_command(ssh, 
                    f"tmux send-keys -t {tmux_session_name}.0 'conda activate {conda_env_name}' ENTER")
        time.sleep(5)
    else:
        print("tmux session was found :")
        print(ssh_stdout)


    """confirm if session is created"""
    _, ssh_stdout, _ = execute_ssh_command(ssh,"tmux ls")
    assert tmux_session_name in ssh_stdout
    print(ssh_stdout)


    """check if vllm server is already running"""
    _,vllm_status = \
        is_running_vllm_server( ssh,  tmux_session_name)
    if vllm_status not in ['running', 'loading']:
        """start vllm server"""
        print( "Vllm server starting... " )
        start_vllm_server = f"python -m vllm.entrypoints.openai.api_server --model {model_id} --port {port}"
        _, _, _ = execute_ssh_command(ssh,
            f"tmux send-keys -t {tmux_session_name}.0 '{start_vllm_server}' ENTER")
        _,vllm_status = \
            is_running_vllm_server( ssh,  tmux_session_name)
        print( f"vllm status {vllm_status}" )
    else:
        print( f"vllm status {vllm_status}" )