birb-scheduler/birb_scheduler/k8s_buzerator.py

#!/usr/bin/env python3
import logging
import time

from kubernetes import config, client
from kubernetes.client.rest import ApiException

config.load_incluster_config()


def ensure_running_pod_on_site(site: str):
    # Configs can be set in Configuration class directly or using helper utility
    logger = logging.getLogger("k8s")

    with client.ApiClient() as api_client:
        # Create an instance of the API class
        api_instance = client.CustomObjectsApi(api_client)
        group = 'types.kubefed.io'  # str | the custom resource's group
        version = 'v1beta1'  # str | the custom resource's version
        namespace = 'birbnetes'  # str | The custom resource's namespace
        plural = 'federateddeployments'  # str | the custom resource's plural name. For TPRs this would be lowercase plural kind.
        name = 'svm-prefilter-service'  # str | the custom object's name

        try:
            api_response = api_instance.get_namespaced_custom_object(group, version, namespace, plural, name)
        except ApiException as e:
            logger.error("Exception when calling CustomObjectsApi->get_namespaced_custom_object: %s\n" % e)
            raise

        patch_needed = False

        override_id_for_this_site = -1
        for i, override in enumerate(api_response['spec']['overrides']):
            if override['clusterName'] == site:
                override_id_for_this_site = i
                break

        if override_id_for_this_site == -1:
            api_response['spec']['overrides'].append(
                {
                    "clusterName": site,
                    "clusterOverrides": [
                        {
                            "path": "/spec/replicas",
                            "value": 1
                        }
                    ]
                }
            )
            patch_needed = True
        else:
            cluster_override_id = -1
            for i, cluster_override in enumerate(
                    api_response['spec']['overrides'][override_id_for_this_site]['clusterOverrides']
            ):
                if cluster_override['path'] == "/spec/replicas":
                    cluster_override_id = i
                    break

            if cluster_override_id == -1:
                api_response['spec']['overrides'][override_id_for_this_site]['clusterOverrides'].append(
                    {"path": "/spec/replicas", "value": 1}
                )
                patch_needed = True
            else:
                current_replicas = api_response['spec']['overrides'][override_id_for_this_site]['clusterOverrides'][cluster_override_id]['value']
                if current_replicas == 0:
                    api_response['spec']['overrides'][override_id_for_this_site]['clusterOverrides'][cluster_override_id]['value'] = 1
                    patch_needed = True

        if patch_needed:
            logger.info("Patching object to fulfill request...")
            api_instance.patch_namespaced_custom_object(group, version, namespace, plural, name, api_response)
            # TODO: Wait for pod to start up
            time.sleep(7)  # close enough
        else:
            logger.debug("No patching needed...")