spot-instance-protection.py

#!/usr/bin/env python3

import asyncio
import requests
import subprocess

lock = asyncio.Lock()

async def get_token():
    token = None
    while True:
        try:
            token = requests.put(
                "http://169.254.169.254/latest/api/token",
                headers={"X-aws-ec2-metadata-token-ttl-seconds": "21600"},
                timeout=5,
            ).text
            break
        except Exception:
            print("Unable to get token")
            await asyncio.sleep(5)
    return token


async def check_for_spot_termination(token):
    try:
        response = requests.get(
            "http://169.254.169.254/latest/meta-data/spot/instance-action",
            headers={"X-aws-ec2-metadata-token": token},
            timeout=5,
        )
        if 199 > response.status_code > 300:
            print("Spot termination has been detected")
            await lock.acquire() # we only want to do this once, never release the lock
            subprocess.run(
                [
                    "/home/skillsadmin/.local/bin/rudderctl",
                    "dump",
                    "push",
                    "--confirm",
                ]
            )
    except Exception:
        pass
    print("No Spot termination detected")


async def loop():
    ctr = 0
    token = await get_token()
    while True:
        asyncio.create_task(check_for_spot_termination(token))
        await asyncio.sleep(5)
        ctr += 1
        if ctr % 4200:
            token = await get_token()


def daemon():
    asyncio.run(loop())


if __name__ == "__main__":
    daemon()