#  Copyright 2019 Amazon.com, Inc. or its affiliates. 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.
#  A copy of the License is located at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  or in the "license" file accompanying this file. This file 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.
import abc
import base64
import json
import logging
import os
import random
import shlex
from dataclasses import dataclass, asdict
from typing import List, Dict
from urllib.parse import urlparse

from bai_kafka_utils.events import DownloadableContent, BenchmarkEvent
from bai_kafka_utils.events import FileSystemObject
from bai_kafka_utils.executors.descriptor import (
    DescriptorError,
    BenchmarkDescriptor,
    DistributedStrategy,
    MPI_JOB_LAUNCHER,
    MPI_JOB_WORKER,
)
from bai_kafka_utils.utils import METRICS_PUSHER_CUSTOM_LABEL_PREFIX
from ruamel import yaml

from executor import SERVICE_NAME
from executor.config import ExecutorConfig
from transpiler.config import BaiConfig, EnvironmentInfo, BaiScriptSource
from transpiler.kubernetes_spec_logic import (
    ConfigTemplate,
    VolumeMount,
    Volume,
    EmptyDirVolumeSource,
    KubernetesRootObjectHelper,
)

SCRIPT_UNPACK_MODIFICATOR = "unpack_in_place"

SCRIPT_PULLER_CHMOD = 777

BENCHMARK_CONTAINER = "benchmark"
DATA_PULLER_CONTAINER = "data-puller"
METRICS_PUSHER_CONTAINER = "metrics-pusher"
METRICS_EXTRACTOR_CONTAINER = "metrics-extractor"
DATASETS_VOLUME_NAME = "datasets-volume"
SCRIPTS_VOLUME_NAME = "scripts-volume"
SCRIPTS_PULLER_CONTAINER = "script-puller"
INFERENCE_SERVER_LOCK_CONTAINER = "inference-server-lock"
INFERENCE_SERVER_CONTAINER = "inference-server"
# To make things easier it's the same path in puller and benchmark
SCRIPTS_MOUNT_PATH = "/bai/scripts"

SHARED_MEMORY_VOLUME = "dshm"
SHARED_MEMORY_VOLUME_MOUNT = "/dev/shm"

BENCHMARK_SERVICE_ACCOUNT = "benchmark"
INFERENCE_SERVER_SERVICE_ACCOUNT = "inference-benchmark"


@dataclass
class PullerDataSource:
    name: str
    puller_path: str


logger = logging.getLogger(__name__)


class BaiKubernetesObjectBuilder(metaclass=abc.ABCMeta):
    def __init__(
        self, job_id: str, descriptor: BenchmarkDescriptor, config: BaiConfig, template_name: str, event: BenchmarkEvent
    ):
        self.job_id = job_id.lower()
        self.inference_server_job_id = f"is-{self.job_id}"
        self.descriptor = descriptor
        self.config_template = ConfigTemplate(self.read_template(template_name))
        self.config = config
        self.event = event

        self.root: KubernetesRootObjectHelper = None

        self.config_template.feed({"descriptor": descriptor})
        self.config_template.feed({"config": config})
        self.config_template.feed({"event": event})
        self.config_template.feed({"service_name": SERVICE_NAME})
        self.config_template.feed({"job_id": self.job_id})
        self.config_template.feed({"inference_server_job_id": self.inference_server_job_id})

    @staticmethod
    def get_metrics_from_descriptor(descriptor: BenchmarkDescriptor):
        metrics = []
        if descriptor.output:
            for metric in descriptor.output.metrics:
                metric.pattern = base64.b64encode(metric.pattern.encode("utf-8")).decode("utf-8")
                metrics.append(asdict(metric))
        return json.dumps(metrics)

    @staticmethod
    def get_server_metrics_from_descriptor(descriptor: BenchmarkDescriptor):
        metrics = []
        if descriptor.is_inference_strategy() and descriptor.server.output:
            for metric in descriptor.server.output.metrics:
                metric.pattern = base64.b64encode(metric.pattern.encode("utf-8")).decode("utf-8")
                metrics.append(asdict(metric))
        return json.dumps(metrics)

    @staticmethod
    def read_template(template_name: str) -> str:
        current_dir = os.path.dirname(os.path.abspath(__file__))
        templates_dir = os.path.join(current_dir, "templates")

        with open(os.path.join(templates_dir, template_name), "r") as f:
            contents = f.read()
        return contents

    @abc.abstractmethod
    def _feed_additional_template_values(self):
        pass

    @abc.abstractmethod
    def _update_root_k8s_object(self):
        pass

    def build(self):
        # if the root object has already been built don't build it again
        if self.root:
            return

        self._feed_additional_template_values()
        self.root = self.config_template.build()

        for label, value in self.descriptor.info.labels.items():
            self.root.add_label(label, value)

        self._update_root_k8s_object()
        return self

    def dump_yaml_string(self):
        self.build()
        return self.root.to_yaml()


class ScheduledBenchmarkKubernetedObjectBuilder(BaiKubernetesObjectBuilder):
    def __init__(self, descriptor: BenchmarkDescriptor, config: BaiConfig, job_id: str, event: BenchmarkEvent):
        super().__init__(job_id, descriptor, config, "cron_job.yaml", event)

    def _feed_additional_template_values(self):
        self.config_template.feed({"event_json_str": json.dumps(self.event.to_json())})

    def _update_root_k8s_object(self):
        pass


class BenchmarkKubernetesObjectBuilder(BaiKubernetesObjectBuilder):
    """
    Adds the logic required from BAI into the Kubernetes root object that represents
    launching a benchmark.
    """

    def __init__(
        self,
        descriptor: BenchmarkDescriptor,
        config: BaiConfig,
        data_sources: List[DownloadableContent],
        scripts: List[BaiScriptSource],
        job_id: str,
        *,
        template_name: str,
        event: BenchmarkEvent,
        environment_info: EnvironmentInfo,
        random_object: random.Random = None,
    ):
        """
        Reads the values from the descriptor file into a settings dictionary
        :param descriptor: Descriptor object with the information from the TOML
        :param random_object: An instance of random.Random [optional].
                              This field exists mostly to facilitate testing so that predictable random data is
                              generated.
        """
        super().__init__(job_id, descriptor, config, template_name=template_name, event=event)

        self.data_sources = data_sources
        self.scripts = scripts
        self.environment_info = environment_info
        self.random_object = random_object
        self.internal_env_vars = {}

    def _feed_additional_template_values(self):
        # Using a random AZ is good enough for now
        availability_zone = self.choose_availability_zone(self.descriptor, self.environment_info, self.random_object)
        self.config_template.feed({"availability_zone": availability_zone})
        self.config_template.feed({"metrics": self.get_metrics_from_descriptor(self.descriptor)})

    def _update_root_k8s_object(self):
        pass

    @staticmethod
    def choose_availability_zone(
        descriptor: BenchmarkDescriptor, environment_info: EnvironmentInfo, random_object: random.Random = None
    ):
        if descriptor.hardware.availability_zone:
            if descriptor.hardware.availability_zone in environment_info.availability_zones.values():
                return descriptor.hardware.availability_zone
            raise DescriptorError(f"Invalid zone name {descriptor.hardware.availability_zone}")
        # Try to find the zone id
        if descriptor.hardware.aws_zone_id:
            if descriptor.hardware.aws_zone_id in environment_info.availability_zones:
                return environment_info.availability_zones[descriptor.hardware.aws_zone_id]
            raise DescriptorError(f"Invalid zone id {descriptor.hardware.aws_zone_id}")

        if random_object is None:
            random_object = random.Random()

        return random_object.choice(list(environment_info.availability_zones.values()))

    def add_metrics_pusher_env_var(self, key: str, value: str, prefix: str, **kwargs):
        env_var_name = prefix.upper() + key.upper()
        try:
            metrics_pusher_container = self.root.find_container(METRICS_PUSHER_CONTAINER, **kwargs)
        except ValueError:
            logger.debug(f"Could not add env var {env_var_name} to metrics pusher sidecar.")
            return
        metrics_pusher_env = metrics_pusher_container.get("env", [])
        metrics_pusher_env.append({"name": env_var_name, "value": value})

    def add_server_env_vars(self):
        if not self.descriptor.is_inference_strategy:
            return

        # TODO: Maybe the k8s namespace should be a parameter to the configuration
        namespace = "default"

        env_vars = {"INFERENCE_SERVER_HOST": f"{self.inference_server_job_id}.{namespace}"}
        ports = self.descriptor.server.env.ports
        env_vars["INFERENCE_SERVER_PORT"] = ports[0]
        for index, port in enumerate(ports[1:]):
            env_vars[f"INFERENCE_SERVER_PORT_{index + 1}"] = port

        self.root.add_env_vars(BENCHMARK_CONTAINER, env_vars)

    def add_benchmark_cmd(self, **kwargs):
        if self.descriptor.hardware.strategy in [DistributedStrategy.SINGLE_NODE, DistributedStrategy.INFERENCE]:
            self.add_benchmark_cmd_to_container()
        elif self.descriptor.hardware.strategy == DistributedStrategy.HOROVOD:
            self.add_benchmark_cmd_to_config_map(**kwargs)
        else:
            raise ValueError("Unsupported configuration in descriptor file")

    def add_data_volume_mounts(self, data_sources, container: str, **kwargs):
        benchmark_container = self.root.find_container(container, **kwargs)
        data_mounts = self._get_data_mounts(data_sources)
        self._update_data_puller(data_mounts, data_sources, **kwargs)
        benchmark_container.volumeMounts.extend(self._get_container_volume_mounts(data_mounts))

    def remove_metrics_sidecars(self, **kwargs):
        self.root.remove_container(METRICS_EXTRACTOR_CONTAINER, **kwargs)
        self.root.remove_container(METRICS_PUSHER_CONTAINER, **kwargs)

    def add_shared_memory(self, **kwargs):
        if self.descriptor.env.extended_shm:
            self._add_extended_shm(**kwargs)

    def _add_extended_shm(self, container=BENCHMARK_CONTAINER, **kwargs):
        pod_spec_volumes = self.root.get_pod_spec(**kwargs).volumes
        container_volume_mounts = self.root.find_container(container, **kwargs).volumeMounts

        shared_memory_vol = SHARED_MEMORY_VOLUME

        shm_vol = Volume(name=shared_memory_vol, emptyDir=EmptyDirVolumeSource(medium="Memory"))
        pod_spec_volumes.append(shm_vol)

        shm_vol_mount = VolumeMount(name=shared_memory_vol, mountPath=SHARED_MEMORY_VOLUME_MOUNT)
        container_volume_mounts.append(shm_vol_mount)

    def dump_yaml_string(self):
        return self.root.to_yaml()

    # TODO: Move this method into the Kubernetes object
    def add_benchmark_cmd_to_container(self, **kwargs):
        """
        Extracts the command and args for the benchmark container, formats and inserts them.
        The command is split into args and passed to the container as a list.
        If a benchmark command was specified, command and args are inserted in the container's command field.
        If only args are provided, they are inserted in the container's args field.
        """
        benchmark_container = self.root.find_container(BENCHMARK_CONTAINER, **kwargs)

        def split_args(command: str) -> List[str]:
            if not command:
                return []
            return shlex.split(command)

        if self.descriptor.ml and self.descriptor.ml.benchmark_code:
            cmd = split_args(self.descriptor.ml.benchmark_code)
            args = split_args(self.descriptor.ml.args)
            benchmark_container.command = cmd + args
            benchmark_container.pop("args", None)

        elif self.descriptor.ml and self.descriptor.ml.args:
            benchmark_container.args = split_args(self.descriptor.ml.args)
            benchmark_container.pop("command", None)

        else:
            benchmark_container.pop("command", None)
            benchmark_container.pop("args", None)

    # TODO: Move this method into the Kubernetes object
    def add_benchmark_cmd_to_config_map(self, **kwargs):
        """
        Adds the benchmark code and args to the entrypoint configmap.
        """
        if self.descriptor.ml and self.descriptor.ml.benchmark_code:
            benchmark_cmd = self.descriptor.ml.benchmark_code + (self.descriptor.ml.args or "")
            # Using yaml.PreservedScalarString so multiline strings are printed properly
            # HACK: Depending on yaml here makes this module depend on knowledge of how our Kubernetes classes are
            # implemented
            # strip whitespaces from lines as they can cause the script to fail
            benchmark_cmd = "\n".join([line.strip() for line in benchmark_cmd.split("\n")])
            benchmark_cmd = yaml.scalarstring.PreservedScalarString(benchmark_cmd)

            entrypoint_config_map = self.root.find_config_map(f"entrypoint-{self.job_id}")
            entrypoint_config_map.data["entrypoint.sh"] = benchmark_cmd

        # If no benchmark_code is specified, we don't need to use the configmap as entrypoint
        else:
            self.add_benchmark_cmd_to_container(**kwargs)

    def _get_data_mounts(self, data_sources: List[DownloadableContent]) -> Dict[str, PullerDataSource]:
        """
        Processes the input data sources to get a dict with the required data volumes
        :param data_sources:
        :return:
        """
        # Data destination paths and the corresponding mounted vols

        data_vols = {}

        for idx, data_source in enumerate(data_sources):
            if data_source.path in data_vols:
                continue
            name = "p" + str(idx)
            puller_path = f"/data/{name}"
            data_vols[data_source.path] = PullerDataSource(name, puller_path)

        return data_vols

    @staticmethod
    def _get_container_volume_mounts(data_volumes: Dict[str, PullerDataSource]) -> List[VolumeMount]:
        vol_mounts = []

        for dest_path, vol in data_volumes.items():
            vol_mounts.append(VolumeMount(name=DATASETS_VOLUME_NAME, mountPath=dest_path, subPath=vol.name))
        return vol_mounts

    @staticmethod
    def _get_s3_bucket_and_object_from_uri(downloadable_content: DownloadableContent):
        parsed_uri = urlparse(downloadable_content.dst)
        scheme = parsed_uri.scheme
        s3_bucket = parsed_uri.netloc
        s3_object = parsed_uri.path[1:]

        if scheme.lower() != "s3":
            raise DescriptorError(
                f"Unexpected scheme in data source src: {scheme}." f" Fetched content is {downloadable_content}"
            )

        return s3_bucket, s3_object

    def _update_data_puller(
        self, data_volumes: Dict[str, PullerDataSource], data_sources: List[DownloadableContent], **kwargs
    ):
        """
        Completes the data puller by adding the required arguments and volume mounts.
        If no data sources are found, the data puller is deleted.
        """
        if not data_sources:
            self.root.remove_container(DATA_PULLER_CONTAINER, **kwargs)
            self.root.remove_volume(DATASETS_VOLUME_NAME, **kwargs)
            return

        data_puller = self.root.find_container(DATA_PULLER_CONTAINER, **kwargs)

        s3_bucket, _ = self._get_s3_bucket_and_object_from_uri(data_sources[0])

        s3_objects = []
        for s in data_sources:
            _, s3_object = self._get_s3_bucket_and_object_from_uri(s)
            s3_objects.append(
                "{object},{chmod},{path_name}".format(
                    object=s3_object, chmod=self.config.puller_mount_chmod, path_name=data_volumes[s.path].name
                )
            )

        puller_args = [s3_bucket, ":".join(s3_objects)]
        data_puller.args = puller_args

    def add_env_vars(self, **kwargs):
        self.root.add_env_vars(BENCHMARK_CONTAINER, self.descriptor.env.vars, **kwargs)
        self.root.add_env_vars(BENCHMARK_CONTAINER, self.internal_env_vars, **kwargs)

    def add_scripts(self, scripts: List[BaiScriptSource], **kwargs):
        if not scripts:
            self.root.remove_container(SCRIPTS_PULLER_CONTAINER, **kwargs)
            self.root.remove_volume(SCRIPTS_VOLUME_NAME, **kwargs)
            return
        script_puller = self.root.find_container(SCRIPTS_PULLER_CONTAINER, **kwargs)
        s3_objects = []

        for inx, s in enumerate(scripts):
            s3_objects.append(
                "{object},{chmod},{path_name},{unpack}".format(
                    object=s.object,
                    chmod=SCRIPT_PULLER_CHMOD,
                    path_name=f"{SCRIPTS_MOUNT_PATH}/s{inx}",
                    unpack=SCRIPT_UNPACK_MODIFICATOR,
                )
            )

        puller_args = [scripts[0].bucket, ":".join(s3_objects)]
        script_puller.args = puller_args

        self.internal_env_vars["BAI_SCRIPTS_PATH"] = SCRIPTS_MOUNT_PATH


class SingleRunBenchmarkKubernetesObjectBuilder(BenchmarkKubernetesObjectBuilder):
    """
    Adds the logic into the Kubernetes root object that represents launching a benchmark for single node job
    """

    def __init__(
        self,
        descriptor: BenchmarkDescriptor,
        config: BaiConfig,
        data_sources: List[DownloadableContent],
        scripts: List[BaiScriptSource],
        job_id: str,
        *,
        template_name: str,
        event: BenchmarkEvent,
        environment_info: EnvironmentInfo,
        random_object: random.Random = None,
    ):
        super().__init__(
            descriptor,
            config,
            data_sources,
            scripts,
            job_id,
            template_name=template_name,
            event=event,
            environment_info=environment_info,
            random_object=random_object,
        )

    def _feed_additional_template_values(self):
        super()._feed_additional_template_values()

    def _update_root_k8s_object(self):
        self.root.set_service_account(BENCHMARK_SERVICE_ACCOUNT)
        self.add_data_volume_mounts(self.data_sources, BENCHMARK_CONTAINER)
        self.add_scripts(self.scripts)
        self.add_shared_memory()
        self.add_env_vars()
        self.add_benchmark_cmd()

        if not self.descriptor.output:
            self.remove_metrics_sidecars()
        else:
            # Add custom labels metrics pusher
            for label, value in self.descriptor.info.labels.items():
                self.add_metrics_pusher_env_var(label, value, prefix=METRICS_PUSHER_CUSTOM_LABEL_PREFIX)
            if self.descriptor.custom_params and self.descriptor.custom_params.dashboard:
                self.root.add_label("dashboard-name", self.descriptor.custom_params.dashboard)
                self.add_metrics_pusher_env_var(
                    "dashboard-name", self.descriptor.custom_params.dashboard, prefix=METRICS_PUSHER_CUSTOM_LABEL_PREFIX
                )
            if self.descriptor.custom_params and self.descriptor.custom_params.region:
                self.root.add_label("region", self.descriptor.custom_params.dashboard)
                self.add_metrics_pusher_env_var(
                    "region", self.descriptor.custom_params.region, prefix=METRICS_PUSHER_CUSTOM_LABEL_PREFIX
                )

        if self.event.parent_action_id:
            self.root.add_label("parent-action-id", self.event.parent_action_id)
            self.add_metrics_pusher_env_var(
                "parent-action-id", self.event.parent_action_id, prefix=METRICS_PUSHER_CUSTOM_LABEL_PREFIX
            )

        if self.config.suppress_job_affinity:
            self.root.remove_affinity()

        if self.descriptor.is_inference_strategy():
            # Use inference server service account
            self.root.set_service_account(INFERENCE_SERVER_SERVICE_ACCOUNT)
            # add server environment variables
            self.add_server_env_vars()

        # remove server lock init container from non-inference strategy
        # single node benchmarks
        elif self.descriptor.hardware.strategy == DistributedStrategy.SINGLE_NODE:
            self.root.remove_container(INFERENCE_SERVER_LOCK_CONTAINER)


class HorovodJobKubernetesObjectBuilder(BenchmarkKubernetesObjectBuilder):
    """
    Adds the logic into the Kubernetes root object that represents launching a benchmark for MPIJob
    """

    def __init__(
        self,
        descriptor: BenchmarkDescriptor,
        config: BaiConfig,
        data_sources: List[DownloadableContent],
        scripts: List[BaiScriptSource],
        job_id: str,
        *,
        template_name: str,
        event: BenchmarkEvent,
        environment_info: EnvironmentInfo,
        random_object: random.Random = None,
    ):
        super().__init__(
            descriptor,
            config,
            data_sources,
            scripts,
            job_id,
            template_name=template_name,
            event=event,
            environment_info=environment_info,
            random_object=random_object,
        )

    def _feed_additional_template_values(self):
        super()._feed_additional_template_values()

    def _update_root_k8s_object(self):
        self.add_data_volume_mounts(self.data_sources, BENCHMARK_CONTAINER, mpiReplicaType=MPI_JOB_WORKER)
        self.add_benchmark_cmd(mpiReplicaType=MPI_JOB_LAUNCHER)

        for replica_type in [MPI_JOB_LAUNCHER, MPI_JOB_WORKER]:
            self.root.set_service_account(BENCHMARK_SERVICE_ACCOUNT, mpiReplicaType=replica_type)
            self.add_scripts(self.scripts, mpiReplicaType=replica_type)
            self.add_shared_memory(mpiReplicaType=replica_type)
            self.add_env_vars(mpiReplicaType=replica_type)

            if not self.descriptor.output:
                self.remove_metrics_sidecars(mpiReplicaType=replica_type)
            else:
                # Add custom labels metrics pusher
                for label, value in self.descriptor.info.labels.items():
                    self.add_metrics_pusher_env_var(
                        label, value, prefix=METRICS_PUSHER_CUSTOM_LABEL_PREFIX, mpiReplicaType=replica_type
                    )
                if self.descriptor.custom_params and self.descriptor.custom_params.dashboard:
                    self.root.add_label("dashboard-name", self.descriptor.custom_params.dashboard)
                    self.add_metrics_pusher_env_var(
                        "dashboard-name",
                        self.descriptor.custom_params.dashboard,
                        prefix=METRICS_PUSHER_CUSTOM_LABEL_PREFIX,
                        mpiReplicaType=replica_type,
                    )
                if self.descriptor.custom_params and self.descriptor.custom_params.region:
                    self.root.add_label("region", self.descriptor.custom_params.dashboard)
                    self.add_metrics_pusher_env_var(
                        "region",
                        self.descriptor.custom_params.region,
                        prefix=METRICS_PUSHER_CUSTOM_LABEL_PREFIX,
                        mpiReplicaType=replica_type,
                    )

            if self.event.parent_action_id:
                self.root.add_label("parent-action-id", self.event.parent_action_id)
                self.add_metrics_pusher_env_var(
                    "parent-action-id",
                    self.event.parent_action_id,
                    prefix=METRICS_PUSHER_CUSTOM_LABEL_PREFIX,
                    mpiReplicaType=replica_type,
                )

            if self.config.suppress_job_affinity:
                self.root.remove_affinity(mpiReplicaType=replica_type)


class InferenceServerJobKubernetedObjectBuilder(BenchmarkKubernetesObjectBuilder):
    def __init__(
        self,
        descriptor: BenchmarkDescriptor,
        config: BaiConfig,
        job_id: str,
        fetched_models: List[DownloadableContent],
        event: BenchmarkEvent,
        environment_info: EnvironmentInfo,
        random_object: random.Random = None,
    ):
        super().__init__(
            descriptor,
            config,
            [],
            [],
            job_id,
            template_name="inference_server_job.yaml",
            event=event,
            environment_info=environment_info,
            random_object=random_object,
        )
        self.fetched_models = fetched_models

    def _feed_additional_template_values(self):
        super()._feed_additional_template_values()
        self.config_template.feed({"server_metrics": self.get_server_metrics_from_descriptor(self.descriptor)})

    def _update_root_k8s_object(self):
        self.root.set_service_account(INFERENCE_SERVER_SERVICE_ACCOUNT)
        self.add_shared_memory()
        self.add_server_command_to_container()
        self.add_data_volume_mounts(self.fetched_models, INFERENCE_SERVER_CONTAINER)
        self.root.add_env_vars(INFERENCE_SERVER_CONTAINER, self.descriptor.server.env.vars)
        self.root.add_tcp_ports_to_service(self.inference_server_job_id, self.descriptor.server.env.ports)
        self.root.add_container_ports_to_container(INFERENCE_SERVER_CONTAINER, self.descriptor.server.env.ports)
        if self.descriptor.server.env.readiness_probe:
            self.root.add_readiness_probe(
                container_name=INFERENCE_SERVER_CONTAINER,
                probe=self.descriptor.server.env.readiness_probe,
                default_port=self.descriptor.server.env.ports[0],
            )
        if self.event.parent_action_id:
            self.root.add_label("parent-action-id", self.event.parent_action_id)
        if not self.descriptor.server.output:
            self.remove_metrics_sidecars()
        else:
            # Add custom labels to metrics pusher as env variables
            for label, value in self.descriptor.info.labels.items():
                self.add_metrics_pusher_env_var(label, value, METRICS_PUSHER_CUSTOM_LABEL_PREFIX)
            if self.descriptor.custom_params and self.descriptor.custom_params.dashboard:
                self.root.add_label("dashboard-name", self.descriptor.custom_params.dashboard)
                self.add_metrics_pusher_env_var(
                    "dashboard-name", self.descriptor.custom_params.dashboard, prefix=METRICS_PUSHER_CUSTOM_LABEL_PREFIX
                )
            if self.descriptor.custom_params and self.descriptor.custom_params.region:
                self.root.add_label("region", self.descriptor.custom_params.dashboard)
                self.add_metrics_pusher_env_var(
                    "region", self.descriptor.custom_params.region, prefix=METRICS_PUSHER_CUSTOM_LABEL_PREFIX
                )
        if self.config.suppress_job_affinity:
            self.root.remove_affinity()

    def add_shared_memory(self):
        if self.descriptor.server.env.extended_shm:
            self._add_extended_shm(container=INFERENCE_SERVER_CONTAINER)

    # TODO: Move this method into the Kubernetes object
    def add_server_command_to_container(self):
        """
        Extracts the command and args for the server container, formats and inserts them.
        The command is split into args and passed to the container as a list.
        If a server command was specified, command and args are inserted in the container's command field.
        If only args are provided, they are inserted in the container's args field.
        """
        server_container = self.root.find_container(INFERENCE_SERVER_CONTAINER)

        def split_args(command: str) -> List[str]:
            if not command:
                return []
            return shlex.split(command)

        cmd = split_args(self.descriptor.server.env.start_command)
        args = split_args(self.descriptor.server.env.start_command_args)
        server_container.command = cmd + args
        server_container.pop("args", None)


def create_scripts(scripts: List[FileSystemObject]) -> List[BaiScriptSource]:
    return list(map(BaiScriptSource, scripts)) if scripts else []


def create_single_run_benchmark_bai_k8s_builder(
    descriptor: BenchmarkDescriptor,
    bai_config: BaiConfig,
    fetched_data_sources: List[DownloadableContent],
    scripts: List[FileSystemObject],
    job_id: str,
    *,
    event: BenchmarkEvent,
    environment_info: EnvironmentInfo,
    extra_bai_config_args: Dict = None,
) -> BaiKubernetesObjectBuilder:
    """
    Builds a BaiKubernetesObjectBuilder object
    :param event: The event that triggered this execution
    :param descriptor: The descriptor.
    :param bai_config: Configuration values.
    :param fetched_data_sources: list of fetched data sources, as generated by the fetcher.
    :param job_id: str
    :param environment_info: Information on the environment that BAI is running on.
    :param extra_bai_config_args: An optional Dict which will be forwarded to the `BaiConfig` object created.
    :return:
    """
    if extra_bai_config_args is None:
        extra_bai_config_args = {}

    distributed_strategy_metadata = {
        DistributedStrategy.SINGLE_NODE: {
            "template_name": "job_single_node.yaml",
            "strategy_type_class": SingleRunBenchmarkKubernetesObjectBuilder,
        },
        DistributedStrategy.INFERENCE: {
            # Benchmark job component of inference strategy
            "template_name": "job_single_node.yaml",
            "strategy_type_class": SingleRunBenchmarkKubernetesObjectBuilder,
        },
        DistributedStrategy.HOROVOD: {
            "template_name": "mpi_job_horovod.yaml",
            "strategy_type_class": HorovodJobKubernetesObjectBuilder,
        },
    }

    template_name = distributed_strategy_metadata.get(descriptor.hardware.strategy)["template_name"]
    if not template_name:
        raise ValueError(f"Unsupported distributed strategy in descriptor file: '{descriptor.hardware.strategy}'")

    bai_scripts = create_scripts(scripts)

    bai_k8s_builder = distributed_strategy_metadata.get(descriptor.hardware.strategy)["strategy_type_class"](
        descriptor,
        bai_config,
        fetched_data_sources,
        bai_scripts,
        job_id,
        template_name=template_name,
        event=event,
        environment_info=environment_info,
        **extra_bai_config_args,
    )

    return bai_k8s_builder.build()


def create_inference_server_bai_k8s_builder(
    descriptor: BenchmarkDescriptor,
    bai_config: BaiConfig,
    job_id: str,
    fetched_models: List[DownloadableContent],
    *,
    event: BenchmarkEvent,
    environment_info: EnvironmentInfo,
    extra_bai_config_args: Dict = None,
) -> BaiKubernetesObjectBuilder:
    """
    Builds a BaiKubernetesObjectBuilder object
    :param event: The event that triggered this execution
    :param descriptor: The descriptor.
    :param bai_config: Configuration values.
    :param job_id: str
    :param environment_info: Information on the environment that BAI is running on
    :param extra_bai_config_args: An optional Dict which will be forwarded to the `BaiConfig` object created
    :return:
    """

    if extra_bai_config_args is None:
        extra_bai_config_args = {}

    bai_k8s_builder = InferenceServerJobKubernetedObjectBuilder(
        descriptor,
        bai_config,
        job_id,
        fetched_models,
        event=event,
        environment_info=environment_info,
        **extra_bai_config_args,
    )
    return bai_k8s_builder.build()


def create_scheduled_benchmark_bai_k8s_builder(
    descriptor: BenchmarkDescriptor, bai_config: BaiConfig, job_id: str, event: BenchmarkEvent
) -> ScheduledBenchmarkKubernetedObjectBuilder:
    """
    Builds a BaiKubernetesObjectBuilder object
    :param event: The event that triggered this execution
    :param descriptor: The descriptor.
    :param bai_config: Configuration values.
    :param job_id: str
    :return:
    """
    bai_k8s_builder = ScheduledBenchmarkKubernetedObjectBuilder(descriptor, bai_config, job_id, event)
    return bai_k8s_builder.build()


def create_job_yaml_spec(
    descriptor_contents: Dict[str, str],
    executor_config: ExecutorConfig,
    fetched_data_sources: List[DownloadableContent],
    fetched_models: List[DownloadableContent],
    scripts: List[FileSystemObject],
    job_id: str,
    *,
    event: BenchmarkEvent,
    extra_bai_config_args=None,
) -> str:
    """
    Creates the YAML spec file corresponding to a descriptor passed as parameter
    :param event: event that triggered this execution
    :param descriptor_contents: dict containing the parsed descriptor
    :param executor_config: configuration for the transpiler
    :param fetched_data_sources: list of fetched data sources, as generated by the fetcher
    :param fetched_models: list of fetched models, as generated by the fetcher
    :param scripts: list of supplied scripts
    :param job_id: str
    :param extra_bai_config_args: An optional Dict which will be forwarded to the `BaiConfig` object created
    :return: Tuple with (yaml string for the given descriptor, job_id)
    """
    descriptor = BenchmarkDescriptor.from_dict(descriptor_contents, executor_config.descriptor_config)

    bai_k8s_benchmark_job_builder = create_single_run_benchmark_bai_k8s_builder(
        descriptor,
        executor_config.bai_config,
        fetched_data_sources,
        scripts,
        job_id,
        event=event,
        environment_info=executor_config.environment_info,
        extra_bai_config_args=extra_bai_config_args,
    )

    if descriptor.hardware.strategy != DistributedStrategy.INFERENCE:
        return bai_k8s_benchmark_job_builder.dump_yaml_string()

    bai_k8s_inference_server_job_builder = create_inference_server_bai_k8s_builder(
        descriptor,
        executor_config.bai_config,
        job_id,
        fetched_models,
        event=event,
        environment_info=executor_config.environment_info,
        extra_bai_config_args=extra_bai_config_args,
    )

    return (
        f"{bai_k8s_benchmark_job_builder.dump_yaml_string()}"
        f"---\n"
        f"{bai_k8s_inference_server_job_builder.dump_yaml_string()}"
    )


def create_scheduled_job_yaml_spec(
    descriptor_contents: Dict, executor_config: ExecutorConfig, job_id: str, event: BenchmarkEvent
) -> str:
    """
    Creates the YAML spec file corresponding to a descriptor passed as parameter
    :param event: event that triggered this execution
    :param descriptor_contents: dict containing the parsed descriptor
    :param executor_config: configuration for the transpiler
    :param job_id: str
    :return: Tuple with (yaml string for the given descriptor, job_id)
    """
    descriptor = BenchmarkDescriptor.from_dict(descriptor_contents, executor_config.descriptor_config)
    bai_k8s_builder = create_scheduled_benchmark_bai_k8s_builder(
        descriptor, executor_config.bai_config, job_id, event=event
    )
    return bai_k8s_builder.dump_yaml_string()
