Approximate an On-Demand count per Ocean VNG

ocean_vng_od_count.py

#!/opt/homebrew/bin/python3
"""
Proof of Concept: On-Demand Count per Ocean VNG (AWS K8s). Need to test how Ocean rounds to nearest whole instance based on Spot Percentage.

Purpose: Currently, Ocean does not support the configuration of 'On-Demand Count', only Elastigroup does. By running this script with AWS Lambda + EventBridge CRON scheduling, on say an hourly basis, this functionality can be approximated.

How it works:
1. Pull the current node count of the Ocean cluster via the API
2. Breakdown node count into separate VNGs
3. For each VNG (or desired VNG), calculate what Risk (Spot Percentage) would equate to having 1x On-Demand instance ( using math.ceil() )
4. For each VNG, Update the configuration via the API

Note: Ensure that continuous optimization is enabled, such that Ocean can use the fix-strategy process to re-adjust the spot percentage.

Further Development: Adjust the API call to allow for functionality with AWS ECS, Azure AKS, and GCP GKE.
"""

try:
    import requests
    import os
    import re
    import json
    import math
except ModuleNotFoundError:
    print("Please ensure proper modules are installed:\npip3 install requests")
    os._exit(1)

# Get Required Variables for Authentication
TOKEN = os.environ.get('SPOTINST_TOKEN')
ACCOUNT = os.environ.get('SPOTINST_ACCOUNT_AWS')

# Check
if not (TOKEN and ACCOUNT):
    print(f"Please set environment variables for token and account id.\n")
    os._exit(1)


class Endpoint:
    def __init__(self, url, body=None):
        """API Endpoint. Account and Token set by default
        
        Arguments:
        url: <class 'str'> representing URL endpoint of API
        body: <class 'str'> valid JSON dumped into a string
        """

        self.url = url
        self.body = body
        self.query_params = dict({"accountId":ACCOUNT})
        self.headers = dict({
            'User-Agent': '_Custom_SDK_Agent',
            'Content-Type': 'application/json',
            'Authorization': 'Bearer ' + TOKEN
        })

    def _GET(self):
        """Send a GET request"""
        print(f"\nFetching node count for {re.findall(ocean_regex, self.url)[0]}...")
        result = requests.get(self.url, params=self.query_params, headers=self.headers)

        if result.status_code == requests.codes.ok:
            data = json.loads(result.content.decode('utf-8'))
            print(f"Success! Fetched node count for {re.findall(ocean_regex, self.url)[0]} -> {result.status_code} {result.reason}")
            return data['response']['items']

        else:
            print(f"Error: {result.reason}\n\n{result.text}\n")
            os._exit(1)

    def _PUT(self):
        """Send a POST request"""
        print(f"\nUpdating Spot Percentage for VNG {re.findall(vng_regex, self.url)[0]}...")
        result = requests.put(self.url, params=self.query_params, data=self.body, headers=self.headers)

        if result.status_code == requests.codes.ok:
            data = json.loads(result.content.decode('utf-8'))
            print(f"Success! Updated Spot Percentage for VNG {re.findall(vng_regex, self.url)[0]} -> {result.status_code} {result.reason}")
            return data['response']['items']

        else:
            print(f"Error: {result.reason}\n\n{result.text}\n")
            os._exit(1)


if __name__ == "__main__":

    # Variables: EDIT
    OD_NODE_COUNT_PER_VNG = 1
    OCEAN_ID = 'o-4b1de2cd'

    # Variables: DO NOT EDIT
    summary = [f"\nComplete! Updated Spot Percentages"]
    ocean_regex = '(o-[0-9a-z]{8})'
    vng_regex = '(ols-[0-9a-z]{8})'

    # Run API Call: Get Current Node Count
    __base_ocean_url = f'https://api.spotinst.io/ocean/aws/k8s/cluster/{OCEAN_ID}/nodes'
    get_ocean_node_count = Endpoint(__base_ocean_url)._GET()

    # Divide Current Node Count into respective VNGs
    vng_node_count = {}
    for node in get_ocean_node_count:
        _vng = node['launchSpecId']

        if _vng not in vng_node_count:
            vng_node_count[_vng] = 1
        else:
            vng_node_count[_vng] += 1

    # Calculate new Spot Ratio
    vng_new_spot_ratios = {}
    for vng, count in vng_node_count.items():
        new_ratio = 100 - math.ceil((OD_NODE_COUNT_PER_VNG * 100) / count)

        # Update Dictionary
        vng_new_spot_ratios[vng] = new_ratio

    # Update Risk: All VNGs on Ocean Cluster
    for vng_id, new_spot_ratio in vng_new_spot_ratios.items():

        # Get New Spot Ratio
        new_spot_ratio = vng_new_spot_ratios.get(vng_id)
        
        # Update VNG Configuration
        __base_vng_url = f'https://api.spotinst.io/ocean/aws/k8s/launchSpec/{vng_id}'
        body = json.dumps({
            "launchSpec": {
                "strategy": {
                "spotPercentage": new_spot_ratio
                }
            }
        })
        update_vng = Endpoint(__base_vng_url, body=body)
        update_vng._PUT()

        # Output: Append VNG row item
        summary.append(f"{vng_id}: {new_spot_ratio}%")

    # Output: Print to terminal
    print('\n'.join(summary))