Automatically Shutting Down and Deallocating an AZURE or AWS VM (when idle AND no one is logged in via SSH)

stop-idle-azure-aws.py

#!/usr/bin/env python

# shell commands being automated.
# w
# aws ec2 stop-instances --instance-id  INSTANCE_ID
# az vm deallocate --name NAME --resource-group RESOURCE_GROUP

"""

The script is the easy part, installing it into the unfriendly(imo) cron system and
making all the permissions and paths are set correctly is another issue altogether.
cron will default to run as a root, so your scripts will fail because it is not
running in the correct python environment, also root might not have the paths for
commands like 'python'. So best to run this in the user space where you are testing
this script already. Here is a cron command to add to the current user's cron than
doing it for root.


sudo crontab -u ec2-user -e

Then add a line like this there. Notice how you have to activate the python project
you are in. Also notice I had to say the whole path to activate. Hopefully after
this it should just work.

*   *   *   *   *   source /usr/local/bin/activate pytorch && python /home/ec2-user/workspace/code/stop_idle_aws.py > /tmp/stop_status

"""

import subprocess
import pandas as pd
import io
import psutil

CPU_IDLE_THRESH    = 10 # Since there might be many machines cores, 100 pct cpu might come out to be 10 pct with 10 cores for eg.
USER_IDLE_THRESH   = 3000 # 50 minutes

# If you are dealing with an aws machine.
# change this to proper value or leave it as None.
INSTANCE_ID    = None 

# if you are dealing with an Azure machine. 
# edit the below variables, or leave it as None. 
NAME           = None 
RESOURCE_GROUP = None 

def get_machine_idle():
    """
    Get the current CPU usage percentage.

    The function uses the psutil library to obtain the current CPU utilization.
    `psutil.cpu_percent()` provides a convenient way to calculate CPU usage.
    This function calls `cpu_percent` with a 1-second interval, which is a
    common practice to get a more accurate reading of CPU usage.

    Returns:
        float: The current CPU usage percentage.
    """
    # Retrieve and return the CPU usage percentage.
    # Interval of 10 second for averaging the CPU usage over that period.
    # This provides a more accurate reading than an instantaneous value.
    print("just before psutil call")
    perc = psutil.cpu_percent(interval=10)
    print("machine cpu average for the past 10 seconds is {}".format(perc))
    return perc < CPU_IDLE_THRESH


def wish_seconds(wish):
    """
    Convert the 'w' command idle time format into seconds.
    Supports formats like: 44.00s, 5:10, 1:28m, 3days, etc.
    """
    if "days" in wish:
        unit1 = int(wish.split("days")[0])
        seconds = unit1 * 86400
    elif "m" in wish:
        unit1, unit2 = wish.split(":")
        print(unit1, unit2)
        unit1 = int(unit1)
        unit2 = int(unit2.rstrip("m"))
        seconds = (unit1 * 3600) + (unit2 * 60)
    elif "s" in wish:
        print("wish", wish)
        seconds = int(float(wish.rstrip("s")))
    else:
        unit1, unit2 = wish.split(":")
        print(unit1, unit2)
        unit1 = int(unit1)
        unit2 = int(unit2)
        seconds = (unit1 * 60) + unit2

    return seconds

def get_user_shortest_idle():
    # Execute the 'w' command and get its output
    output = subprocess.check_output(['w'], text=True).strip()  # '-h' to skip header

    # Use pandas to read the output into a DataFrame
    df = pd.read_csv(io.StringIO(output), delim_whitespace=True, skiprows=1)

    # Extract the IDLE times
    idle_times = df['IDLE']

    # Convert the idle times into seconds
    seconds_values = idle_times.apply(wish_seconds)

    # Find the shortest idle time
    shortest = seconds_values.min()

    # if shortest is nan then there is no idle time, so set it to a large value
    # check if shortest is nan
    if shortest != shortest:
        shortest = USER_IDLE_THRESH + 1

    # print the shortest idle time
    print("shortest: ", shortest)

    return shortest > USER_IDLE_THRESH


def main():
    # check if there are any active ssh connections to the vm.
    # this is to prevent the machine shutting down when I am working
    # via the shell.

    print("Getting user idle..")
    user_idle   = get_user_shortest_idle()
    print("Getting mac idle...")
    mac_idle    = get_machine_idle()

    print("user_idle: ", user_idle)
    print("mac_idle: ", mac_idle)

    """
    Check if both the machine and the user are idle, if so, shut down the vm.
    Else print an appropriate message, and leave the vm alone.
    """

    if user_idle:
        if mac_idle:
            print ("Both Users and Machine Seems Idle.., Shutting down!")
            if NAME: # Check if azure machine. 
                subprocess.run(['az', 'vm', 'deallocate', "--name", NAME, "--resource-group", RESOURCE_GROUP])
            elif INSTANCE_ID: # Check if aws machine. 
                subprocess.run(['aws', 'ec2', 'stop-instances', '--instance-ids', INSTANCE_ID])
            else:
                print("Edit NAME if this is an Azure machine, or edit INSTANCE_ID if this is an AWS machine")

            print("Users are idle, but machine is busy, leaving the vm alone..")
    else:
        print ("Users are active, leaving the vm alone..")

main()

@chinniprakash The code solves for this scenario, kindly refer to the function get_logged_in() which is exactly what you want. (provided any of those logins are using ssh to connect. In case they don't you will have to mod this function)

script modified to introduce aws and azure into same script, edit depending on what you are running it on. also more resilient conditions introduced.. the previous ones were a bit flaky under certain circumstances.