community/modules/scripts/htcondor-install/files/autoscaler.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-

# Copyright 2018 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Script for resizing managed instance group (MIG) cluster size based
# on the number of jobs in the Condor Queue.

from absl import app
from absl import flags
from collections import OrderedDict
from datetime import datetime
from pprint import pprint
from googleapiclient import discovery
from oauth2client.client import GoogleCredentials

import argparse
import os
import math
import time
import htcondor
import classad

parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("--p", required=True, help="Project id", type=str)
parser.add_argument(
    "--z",
    required=True,
    help="Name of GCP zone where the managed instance group is located",
    type=str,
)
parser.add_argument(
    "--r",
    required=True,
    help="Name of GCP region where the managed instance group is located",
    type=str,
)
parser.add_argument(
    "--mz",
    required=False,
    help="Enabled multizone (regional) managed instance group",
    action="store_true",
)
parser.add_argument(
    "--g", required=True, help="Name of the managed instance group", type=str
)
parser.add_argument(
    "--i",
    default=0,
    help="Minimum number of idle compute instances",
    type=int
)
parser.add_argument(
    "--c", required=True, help="Maximum number of compute instances", type=int
)
parser.add_argument(
    "--v",
    default=0,
    help="Increase output verbosity. 1-show basic debug info. 2-show detail debug info",
    type=int,
    choices=[0, 1, 2],
)
parser.add_argument(
    "--d",
    default=0,
    help="Dry Run, default=0, if 1, then no scaling actions",
    type=int,
    choices=[0, 1],
)

args = parser.parse_args()

class AutoScaler:
    def __init__(self, multizone=False):

        self.multizone = multizone
        # Obtain credentials
        self.credentials = GoogleCredentials.get_application_default()
        self.service = discovery.build("compute", "v1", credentials=self.credentials)

        if self.multizone:
            self.instanceGroupManagers = self.service.regionInstanceGroupManagers()
        else:
            self.instanceGroupManagers = self.service.instanceGroupManagers()

    # Remove specified instances from MIG and decrease MIG size
    def deleteFromMig(self, node_self_links):
        requestDelInstance = self.instanceGroupManagers.deleteInstances(
            project=self.project,
            **self.zoneargs,
            instanceGroupManager=self.instance_group_manager,
            body={ "instances": node_self_links },
        )

        # execute if not a dry-run
        if not self.dryrun:
            response = requestDelInstance.execute()
            if self.debug > 0:
                pprint(response)
            return response
        return "Dry Run"

    def getInstanceTemplateInfo(self):
        requestTemplateName = self.instanceGroupManagers.get(
            project=self.project,
            **self.zoneargs,
            instanceGroupManager=self.instance_group_manager,
            fields="instanceTemplate",
        )
        responseTemplateName = requestTemplateName.execute()
        template_name = ""

        if self.debug > 1:
            print("Request for the template name")
            pprint(responseTemplateName)

        if len(responseTemplateName) > 0:
            template_url = responseTemplateName.get("instanceTemplate")
            template_url_partitioned = template_url.split("/")
            template_name = template_url_partitioned[len(template_url_partitioned) - 1]

        requestInstanceTemplate = self.service.instanceTemplates().get(
            project=self.project, instanceTemplate=template_name, fields="properties"
        )
        responseInstanceTemplateInfo = requestInstanceTemplate.execute()

        if self.debug > 1:
            print("Template information")
            pprint(responseInstanceTemplateInfo["properties"])

        machine_type = responseInstanceTemplateInfo["properties"]["machineType"]
        is_spot = responseInstanceTemplateInfo["properties"]["scheduling"][
            "preemptible"
        ]
        if self.debug > 0:
            print("Machine Type: " + machine_type)
            print("Is spot: " + str(is_spot))
        request = self.service.machineTypes().get(
            project=self.project, zone=self.zone, machineType=machine_type
        )
        response = request.execute()
        guest_cpus = response["guestCpus"]
        if self.debug > 1:
            print("Machine information")
            pprint(responseInstanceTemplateInfo["properties"])
        if self.debug > 0:
            print("Guest CPUs: " + str(guest_cpus))

        instanceTemplateInfo = {
            "machine_type": machine_type,
            "is_spot": is_spot,
            "guest_cpus": guest_cpus,
        }
        return instanceTemplateInfo

    def scale(self):
        # diagnosis
        if self.debug > 1:
            print("Launching autoscaler.py with the following arguments:")
            print("project_id: " + self.project)
            print("zone: " + self.zone)
            print("region: " + self.region)
            print(f"multizone: {self.multizone}")
            print("group_manager: " + self.instance_group_manager)
            print("computeinstancelimit: " + str(self.compute_instance_limit))
            print("debuglevel: " + str(self.debug))

        if self.multizone:
            self.zoneargs = {"region": self.region}
        else:
            self.zoneargs = {"zone": self.zone}

        # Each HTCondor scheduler (SchedD), maintains a list of jobs under its
        # stewardship. A full list of Job ClassAd attributes can be found at
        # https://htcondor.readthedocs.io/en/latest/classad-attributes/job-classad-attributes.html
        schedd = htcondor.Schedd()
        # encourage the job queue to start a new negotiation cycle; there are
        # internal unconfigurable rate limits so not guaranteed; this is not
        # strictly required for success, but may reduce latency of autoscaling
        schedd.reschedule()
        REQUEST_CPUS_ATTRIBUTE = "RequestCpus"
        REQUEST_GPUS_ATTRIBUTE = "RequestGpus"
        REQUEST_MEMORY_ATTRIBUTE = "RequestMemory"
        job_attributes = [
            REQUEST_CPUS_ATTRIBUTE,
            REQUEST_GPUS_ATTRIBUTE,
            REQUEST_MEMORY_ATTRIBUTE,
        ]

        instanceTemplateInfo = self.getInstanceTemplateInfo()
        self.is_spot = instanceTemplateInfo["is_spot"]
        self.cores_per_node = instanceTemplateInfo["guest_cpus"]
        print(f"MIG is configured for Spot pricing: {self.is_spot}")
        print("Number of CPU per compute node: " + str(self.cores_per_node))

        # this query will constrain the search for jobs to those that either
        # require spot VMs or do not require Spot VMs based on whether the
        # VM instance template is configured for Spot pricing
        spot_query = classad.ExprTree(f"RequireId == \"{self.instance_group_manager}\"")

        # For purpose of scaling a Managed Instance Group, count only jobs that
        # are idle and likely participated in a negotiation cycle (there does
        # not appear to be a single classad attribute for this).
        # https://htcondor.readthedocs.io/en/latest/classad-attributes/job-classad-attributes.html#JobStatus
        LAST_CYCLE_ATTRIBUTE = "LastNegotiationCycleTime0"
        coll = htcondor.Collector()
        negotiator_ad = coll.query(htcondor.AdTypes.Negotiator, projection=[LAST_CYCLE_ATTRIBUTE])
        if len(negotiator_ad) != 1:
            print(f"There should be exactly 1 negotiator in the pool. There is {len(negotiator_ad)}")
            exit()
        last_negotiation_cycle_time = negotiator_ad[0].get(LAST_CYCLE_ATTRIBUTE)
        if not last_negotiation_cycle_time:
            print(f"The negotiator has not yet started a match cycle. Exiting auto-scaling.")
            exit()

        print(f"Last negotiation cycle occurred at: {datetime.fromtimestamp(last_negotiation_cycle_time)}")
        idle_job_query = classad.ExprTree(f"JobStatus == 1 && QDate < {last_negotiation_cycle_time}")
        idle_job_ads = schedd.query(constraint=idle_job_query.and_(spot_query),
                                    projection=job_attributes)

        total_idle_request_cpus = sum(j[REQUEST_CPUS_ATTRIBUTE] for j in idle_job_ads)
        print(f"Total CPUs requested by idle jobs: {total_idle_request_cpus}")

        if self.debug > 1:
            print("Information about the compute instance template")
            pprint(instanceTemplateInfo)

        # Calculate the minimum number of instances that, for fully packed
        # execute points, could satisfy current job queue
        min_hosts_for_idle_jobs = math.ceil(total_idle_request_cpus / self.cores_per_node)
        if self.debug > 0:
            print(f"Minimum hosts needed: {total_idle_request_cpus} / {self.cores_per_node} = {min_hosts_for_idle_jobs}")

        # Get current number of instances in the MIG
        requestGroupInfo = self.instanceGroupManagers.get(
            project=self.project,
            **self.zoneargs,
            instanceGroupManager=self.instance_group_manager,
        )
        responseGroupInfo = requestGroupInfo.execute()
        current_target = responseGroupInfo["targetSize"]
        print(f"Current MIG target size: {current_target}")

        # Find instances that are being modified by the MIG (currentAction is
        # any value other than "NONE"). A common reason an instance is modified
        # is it because it has failed a health check.
        reqModifyingInstances = self.instanceGroupManagers.listManagedInstances(
            project=self.project,
            **self.zoneargs,
            instanceGroupManager=self.instance_group_manager,
            filter="currentAction != \"NONE\"",
            orderBy="creationTimestamp desc"
        )
        respModifyingInstances = reqModifyingInstances.execute()

        # Find VMs that are idle (no dynamic slots created from partitionable
        # slots) in the MIG handled by this autoscaler
        filter_idle_vms = classad.ExprTree(f"PartitionableSlot && NumDynamicSlots==0")
        filter_claimed_vms = classad.ExprTree(f"PartitionableSlot && NumDynamicSlots>0")
        filter_mig = classad.ExprTree(f"regexp(\".*/{self.instance_group_manager}$\", CloudCreatedBy)")
        # A full list of Machine (StartD) ClassAd attributes can be found at
        # https://htcondor.readthedocs.io/en/latest/classad-attributes/machine-classad-attributes.html
        idle_node_ads = coll.query(htcondor.AdTypes.Startd,
            constraint=filter_idle_vms.and_(filter_mig),
            projection=["Machine", "CloudZone"])

        NODENAME_ATTRIBUTE = "Machine"
        claimed_node_ads = coll.query(htcondor.AdTypes.Startd,
            constraint=filter_claimed_vms.and_(filter_mig),
            projection=[NODENAME_ATTRIBUTE])
        claimed_nodes = [ ad[NODENAME_ATTRIBUTE].split(".")[0] for ad in claimed_node_ads]

        # treat OrderedDict as a set by ignoring key values; this set will
        # contain VMs we would consider deleting, in inverse order of
        # their readiness to join pool (creating, unhealthy, healthy+idle)
        idle_nodes = OrderedDict()
        try:
            modifyingInstances = respModifyingInstances["managedInstances"]
        except KeyError:
            modifyingInstances = []

        print(f"There are {len(modifyingInstances)} VMs being modified by the managed instance group")

        # there is potential for nodes in MIG health check "VERIFYING" state
        # to have already joined the pool and be running jobs
        for instance in modifyingInstances:
            self_link = instance["instance"]
            node_name = self_link.rsplit("/", 1)[-1]
            if node_name not in claimed_nodes:
                idle_nodes[self_link] = "modifying"

        for ad in idle_node_ads:
            node = ad["Machine"].split(".")[0]
            zone = ad["CloudZone"]
            self_link = "https://www.googleapis.com/compute/v1/projects/" + \
                self.project + "/zones/" + zone + "/instances/" + node
            # there is potential for nodes in MIG health check "VERIFYING" state
            # to have already joined the pool and be idle; delete them last
            if self_link in idle_nodes:
                idle_nodes.move_to_end(self_link)
            idle_nodes[self_link] = "idle"
        n_idle = len(idle_nodes)

        print(f"There are {n_idle} VMs being modified or idle in the pool")
        if self.debug > 1:
            print("Listing idle nodes:")
            pprint(idle_nodes)

        # always keep size tending toward the minimum idle VMs requested
        new_target = current_target + self.compute_instance_min_idle - n_idle + min_hosts_for_idle_jobs
        if new_target > self.compute_instance_limit:
            self.size = self.compute_instance_limit
            print(f"MIG target size will be limited by {self.compute_instance_limit}")
        else:
            self.size = new_target

        print(f"New MIG target size: {self.size}")

        if self.debug > 1:
            print("MIG Information:")
            print(responseGroupInfo)

        if self.size == current_target:
            if current_target == 0:
                print("Queue is empty")
            print("Running correct number of VMs to handle queue")
            exit()

        if self.size < current_target:
            print("Scaling down. Looking for nodes that can be shut down")

            if self.debug > 1:
                print("Compute node busy status:")
                for node in idle_nodes:
                    print(node)

            # Shut down idle nodes up to our calculated limit
            nodes_to_delete = list(idle_nodes.keys())[0:current_target-self.size]
            for node in nodes_to_delete:
                print(f"Attempting to delete: {node.rsplit('/',1)[-1]}")
            respDel = self.deleteFromMig(nodes_to_delete)

            if self.debug > 1:
                print("Scaling down complete")

        if self.size > current_target:
            print(
                "Scaling up. Need to increase number of instances to " + str(self.size)
            )
            # Request to resize
            request = self.instanceGroupManagers.resize(
                project=self.project,
                **self.zoneargs,
                instanceGroupManager=self.instance_group_manager,
                size=self.size,
            )
            response = request.execute()
            if self.debug > 1:
                print("Requesting to increase MIG size")
                pprint(response)
                print("Scaling up complete")


def main():

    scaler = AutoScaler(args.mz)

    # Project ID
    scaler.project = args.p  # Ex:'slurm-var-demo'

    # Name of the zone where the managed instance group is located
    scaler.zone = args.z  # Ex: 'us-central1-f'

    # Name of the region where the managed instance group is located
    scaler.region = args.r  # Ex: 'us-central1'

    # The name of the managed instance group.
    scaler.instance_group_manager = args.g  # Ex: 'condor-compute-igm'

    # Default number of cores per instance, will be replaced with actual value
    scaler.cores_per_node = 4

    # Default number of running instances that the managed instance group should maintain at any given time. This number will go up and down based on the load (number of jobs in the queue)
    scaler.size = 0

    scaler.compute_instance_min_idle = args.i

    # Dry run: : 0, run scaling; 1, only provide info.
    scaler.dryrun = args.d > 0

    # Debug level: 1-print debug information, 2 - print detail debug information
    scaler.debug = 0
    if args.v:
        scaler.debug = args.v

    # Limit for the maximum number of compute instance. If zero (default setting), no limit will be enforced by the  script
    scaler.compute_instance_limit = 0
    if args.c:
        scaler.compute_instance_limit = abs(args.c)

    scaler.scale()


if __name__ == "__main__":
    main()