def make_patches()

in tpu-provisioner/admission_controller/admission_controller.py [0:0]


def make_patches(object_in: dict) -> JsonPatch:
    '''Generates a JsonPatch for Job mutations that are based on environment variables.'''
    job_name: str = object_in["metadata"]["name"]
    job_namespace: str = object_in["metadata"]["namespace"]
    modified_object: dict = deepcopy(object_in)

    if "nodeSelector" not in modified_object["spec"]["template"]["spec"]:
        modified_object["spec"]["template"]["spec"]["nodeSelector"] = {}

    # Add job-key node selector unconditionally.
    modified_object["spec"]["template"]["spec"]["nodeSelector"][job_key_label] = job_key_value(job_name, job_namespace)
    webhook_logger.info(f'Job: {job_name} Added nodeSelector: {job_key_label}: {job_key_value(job_name, job_namespace)}')

    if os.environ.get(FORCE_ON_DEMAND) == "true":
        # Remove reservation label if FORCE_ON_DEMAND is set.
        if reservation_name_label in modified_object["spec"]["template"]["spec"]["nodeSelector"]:
            del modified_object["spec"]["template"]["spec"]["nodeSelector"][reservation_name_label]
            webhook_logger.info(f'Job: {job_name} Removed nodeSelector for node label: {reservation_name_label}')
        # Remove spot label if FORCE_ON_DEMAND is set.
        if gke_spot_label in modified_object["spec"]["template"]["spec"]["nodeSelector"]:
            del modified_object["spec"]["template"]["spec"]["nodeSelector"][gke_spot_label]
            webhook_logger.info(f'Job: {job_name} Removed nodeSelector for node label: {gke_spot_label}')

    # Set location hint nodeSelector if RESERVATION_LOCATION_HINT is set.
    location_hint_value: str = os.environ.get(LOCATION_HINT, "")
    if location_hint_value != "":
        modified_object["spec"]["template"]["spec"]["nodeSelector"][gke_location_hint_label] = location_hint_value
        webhook_logger.info(f'Job: {job_name} Added nodeSelector: {gke_location_hint_label}: {location_hint_value}')

    patch: JsonPatch = JsonPatch.from_diff(object_in, modified_object)
    return patch