def AlignByTime(data):
    keyed_data = data | "Flooring time" >> beam.Map(_AssignUniqueIDAndFlooredTimeAsKey)
    five_minute_groups = keyed_data | "Group Data by Keys" >> beam.GroupByKey()
    max_record = five_minute_groups | "Pick Max Record in 5 Minutes" >> beam.ParDo(
        _PickMaxRecord()
    )
    simulated_sample = max_record | "Change VMSample to SimulatedSammple" >> beam.Map(
        _VMSampleToSimulatedSample
    )
    return simulated_sample def SetAbstractMetrics(data, configs):
    metric = configs.metric.WhichOneof("metric")

    if metric == None:
        samples_with_abstract_metric = data
    else:
        if metric in [
            "avg_memory_usage",
            "max_memory_usage",
            "random_sample_memory_usage",
            "assigned_memory",
        ]:
            samples_with_abstract_metric = (
                data
                | "Setting Memory Metrics and Limit"
                >> beam.Map(SetMemoryMetricAndLimit, metric)
            )

        else:
            percentile = None
            if metric == "cpu_usage_percentile":
                percentile = configs.metric.cpu_usage_percentile
            samples_with_abstract_metric = (
                data
                | "Setting Memory Metrics and Limit"
                >> beam.Map(SetCPUMetricAndLimit, metric, percentile)
            )

    return samples_with_abstract_metric def FilterVMSample(data, configs):

    if configs.filter.HasField("start_time") and configs.filter.HasField("end_time"):
        vms_in_time_range = data | "Keep VMs within time range" >> beam.Filter(
            VMsInTimeRange, configs.filter.start_time, configs.filter.end_time
        )
    else:
        vms_in_time_range = data

    if configs.filter.remove_non_top_level_vms == True:
        top_level_vms = vms_in_time_range | "Remove Non Top Level VMs" >> beam.Filter(
            RemoveNonTopLevel
        )
    else:
        top_level_vms = vms_in_time_range

    if configs.filter.priority_range.HasField(
        "lower_bound"
    ) and configs.filter.priority_range.HasField("upper_bound"):
        vms_in_priority_range = (
            top_level_vms
            | "Keep VMs within priority range"
            >> beam.Filter(
                VMsInPriorityRange,
                configs.filter.priority_range.lower_bound,
                configs.filter.priority_range.upper_bound,
            )
        )
    else:
        vms_in_priority_range = top_level_vms

    if configs.filter.scheduling_class_range.HasField(
        "lower_bound"
    ) and configs.filter.scheduling_class_range.HasField("upper_bound"):
        vms_in_scheduling_class_range = (
            vms_in_priority_range
            | "Keep VMs within scheduling class range"
            >> beam.Filter(
                VMsInSchedulingClassRange,
                configs.filter.scheduling_class_range.lower_bound,
                configs.filter.scheduling_class_range.upper_bound,
            )
        )
    else:
        vms_in_scheduling_class_range = vms_in_priority_range

    return vms_in_scheduling_class_range def ResetAndShiftSimulatedTime(data, configs):

    keyed_data = data | "Assign Unique ID as the Key" >> beam.Map(_AssignUniqueIDAsKey)
    per_VM_groups = keyed_data | "Group Data by Unique ID" >> beam.GroupByKey()
    sorted_data = per_VM_groups | "Sort by Simulated Time before Reset" >> beam.ParDo(
        _SortBySimulatedTime()
    )

    lower_bound = configs.reset_and_shift.random_shift.lower_bound
    upper_bound = configs.reset_and_shift.random_shift.upper_bound
    seed = (
        configs.reset_and_shift.seed
        if configs.reset_and_shift.HasField("seed")
        else None
    )

    shifted_data = sorted_data | "Reset and Shift" >> beam.Map(
        _ResetAndShift, lower_bound, upper_bound, seed
    )
    unpacked_groups = shifted_data | "Ungroup the data" >> beam.FlatMap(
        lambda elements: elements
    )
    return unpacked_groups def SetScheduler(data, configs):
    scheduler = configs.scheduler.WhichOneof("scheduler")

    if scheduler == "by_machine_id":
        scheduled_samples = data | "Scheduling by Machine ID" >> beam.Map(
            _ScheduleByMachineID
        )

    if scheduler == "by_vm_unique_id":
        scheduled_samples = data | "Scheduling by VM ID" >> beam.Map(_ScheduleByVMID)

    if scheduler == "at_random":
        if configs.scheduler.at_random.seed != None:
            random.seed(configs.scheduler.at_random.seed)
        scheduled_samples = data | "Scheduling at Random" >> beam.Map(
            _ScheduleAtRandom, configs.scheduler.at_random.num_machines,
        )

    return data def CallFortuneTellerRunner(data, config):

    keyed_data = data | "Assign Simulated Machine ID to Data" >> beam.Map(
        _AssignSimulatedMachineIDAsKey
    )

    grouped_data = (
        keyed_data | "Group by Machine ID before Sorting" >> beam.GroupByKey()
    )

    sorted_data = grouped_data | "Sort by Simulated Time" >> beam.Map(
        _SortBySimulatedTime
    )

    results_schema_without_samples_file = open(
        "simulator/results_schema_without_samples.json"
    )
    results_schema_without_samples = json.load(results_schema_without_samples_file)
    for fortune_teller_config in config.fortune_teller:

        simulation_results = sorted_data | "Calling Fortune Teller Runner For {}".format(
            fortune_teller_config.name
        ) >> beam.Map(
            _FortuneTellerRunner, _FortuneTeller(fortune_teller_config)
        )

        unpacked_simulation_results = simulation_results | "Ungroup the simulation resultsFor {}".format(
            fortune_teller_config.name
        ) >> beam.FlatMap(
            lambda elements: elements
        )

        if config.simulation_result.dataset:
            simulation_result_dataset = config.simulation_result.dataset
            simulation_result_table = (
                config.simulation_result.table
                if config.simulation_result.HasField("table")
                else fortune_teller_config.name
            )

            if fortune_teller_config.save_samples == True:
                # TODO: add support for saving samples
                assert False, "Simulator does not support saving samples at present."
            else:
                unpacked_simulation_results | "Save {} results to BQ table".format(
                    fortune_teller_config.name
                ) >> beam.io.WriteToBigQuery(
                    "{}.{}".format(simulation_result_dataset, simulation_result_table,),
                    schema=results_schema_without_samples,
                    write_disposition=beam.io.BigQueryDisposition.WRITE_TRUNCATE,
                    create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
                )
        else:
            return unpacked_simulation_results class PredictorFactory(object):
    __instance = None

    def __new__(cls):
        if PredictorFactory.__instance is None:
            PredictorFactory.__instance = object.__new__(cls)
            PredictorFactory.__instance.decorator_factories = {}
            PredictorFactory.__instance.concrete_predictor_factories = {}
        return PredictorFactory.__instance

    def RegisterDecorator(self, name, factory_func):
        self.decorator_factories[name] = factory_func

    def RegisterPredictor(self, name, factory_func):
        self.concrete_predictor_factories[name] = factory_func

    def CreatePredictor(self, config):
        name = config.WhichOneof("predictor")

        if name in self.concrete_predictor_factories:
            return self.concrete_predictor_factories[name](getattr(config, name))

        elif name in self.decorator_factories:
            decorated_predictors = []
            for decorated_predictor_config in config.decorated_predictors:
                predictor = self.CreatePredictor(decorated_predictor_config)
                decorated_predictors.append(predictor)
            return self.decorator_factories[name](
                getattr(config, name), decorated_predictors
            )

        else:
            assert False, "Requested predictor or decorator not registered. " class AvgPredictor(StatefulPredictor):
    def __init__(self, config):
        super().__init__(config)
        self.cap_to_limit = config.cap_to_limit

    def CreateState(self, vm_info):
        return _State()

    def UpdateState(self, vm_measure, vm_state):
        limit = vm_measure["sample"]["abstract_metrics"]["limit"]
        usage = vm_measure["sample"]["abstract_metrics"]["usage"]
        if self.cap_to_limit == True:
            usage = min(usage, limit)
        vm_state.total_usage += usage

    def Predict(self, vm_states_and_num_samples):

        vms_avgs = []
        for vm_state_and_num_sample in vm_states_and_num_samples:
            vms_avgs.append(
                vm_state_and_num_sample.vm_state.total_usage
                / vm_state_and_num_sample.vm_num_samples
            )

        predicted_peak = sum(vms_avgs)

        return predicted_peak class MaxPredictor(StatefulPredictor):
    def __init__(self, config):
        super().__init__(config)
        self.cap_to_limit = config.cap_to_limit

    def CreateState(self, vm_info):
        return _State()

    def UpdateState(self, vm_measure, vm_state):
        limit = vm_measure["sample"]["abstract_metrics"]["limit"]
        usage = vm_measure["sample"]["abstract_metrics"]["usage"]
        if self.cap_to_limit == True:
            usage = min(usage, limit)
        vm_state.usage = max(usage, vm_state.usage)

    def Predict(self, vm_states_and_num_samples):

        vms_peaks = []
        for vm_state_and_num_sample in vm_states_and_num_samples:
            vms_peaks.append(vm_state_and_num_sample.vm_state.usage)

        predicted_peak = sum(vms_peaks)

        return predicted_peak class PerVMPercentilePredictor(StatefulPredictor):
    def __init__(self, config):
        super().__init__(config)
        self.num_history_samples = config.num_history_samples
        self.percentile = min(config.percentile, 100)
        self.cap_to_limit = config.cap_to_limit

    def CreateState(self, vm_info):
        return _State(self.num_history_samples)

    def UpdateState(self, vm_measure, vm_state):
        limit = vm_measure["sample"]["abstract_metrics"]["limit"]
        usage = vm_measure["sample"]["abstract_metrics"]["usage"]
        if self.cap_to_limit == True:
            usage = min(usage, limit)
        vm_state.usage.append(usage)
        vm_state.limit.append(limit)

    def Predict(self, vm_states_and_num_samples):

        vms_percentiles = []
        for vm_state_and_num_sample in vm_states_and_num_samples:
            normalized_usage = [
                usage / limit
                for usage, limit in zip(
                    list(vm_state_and_num_sample.vm_state.usage),
                    list(vm_state_and_num_sample.vm_state.limit),
                )
            ]

            vms_percentiles.append(
                np.percentile(np.array(normalized_usage), self.percentile)
                * vm_state_and_num_sample.vm_state.limit[-1]
            )

        predicted_peak = np.sum(vms_percentiles)

        return predicted_peak class PerMachinePercentilePredictor(StatefulPredictor):
    def __init__(self, config):
        super().__init__(config)
        self.num_history_samples = config.num_history_samples
        self.cap_to_limit = config.cap_to_limit
        self.percentile = min(config.percentile, 100)

    def CreateState(self, vm_info):
        return _State(self.num_history_samples)

    def UpdateState(self, vm_measure, vm_state):
        limit = vm_measure["sample"]["abstract_metrics"]["limit"]
        usage = vm_measure["sample"]["abstract_metrics"]["usage"]
        if self.cap_to_limit == True:
            usage = min(usage, limit)
        vm_state.usage.appendleft(usage)
        vm_state.limit.appendleft(limit)

    def Predict(self, vm_states_and_num_samples):

        total_normalized_usage = []
        for idx in range(self.num_history_samples):
            total_limit = 0
            total_usage = 0
            num_vms = 0
            for vm_state_and_num_sample in vm_states_and_num_samples:
                if idx < len(vm_state_and_num_sample.vm_state.usage):
                    total_usage += vm_state_and_num_sample.vm_state.usage[idx]
                    total_limit += vm_state_and_num_sample.vm_state.limit[idx]
                    num_vms += 1
            usage_to_limit_ratio = 0
            if num_vms == 0:
                continue
            if total_limit > 0:
                usage_to_limit_ratio = total_usage / total_limit
            total_normalized_usage.append(usage_to_limit_ratio)

        predicted_peak = min(
            1.0, np.percentile(np.array(total_normalized_usage), self.percentile)
        )
        current_total_limit = 0
        for vm_state_and_num_sample in vm_states_and_num_samples:
            if len(vm_state_and_num_sample.vm_state.limit) > 0:
                current_total_limit += vm_state_and_num_sample.vm_state.limit[0]
        return predicted_peak * current_total_limit class NSigmaPredictor(StatefulPredictor):
    def __init__(self, config):
        super().__init__(config)
        self.num_history_samples = config.num_history_samples
        self.cap_to_limit = config.cap_to_limit
        self.n = config.n

    def CreateState(self, vm_info):
        return _State(self.num_history_samples)

    def UpdateState(self, vm_measure, vm_state):
        limit = vm_measure["sample"]["abstract_metrics"]["limit"]
        usage = vm_measure["sample"]["abstract_metrics"]["usage"]
        if self.cap_to_limit == True:
            usage = min(usage, limit)
        vm_state.usage.appendleft(usage)
        vm_state.limit.appendleft(limit)

    def Predict(self, vm_states_and_num_samples):

        total_normalized_usage = []
        for idx in range(self.num_history_samples):
            total_limit = 0
            total_usage = 0
            num_vms = 0
            for vm_state_and_num_sample in vm_states_and_num_samples:
                if idx < len(vm_state_and_num_sample.vm_state.usage):
                    total_usage += vm_state_and_num_sample.vm_state.usage[idx]
                    total_limit += vm_state_and_num_sample.vm_state.limit[idx]
                    num_vms += 1
            usage_to_limit_ratio = 0
            if num_vms == 0:
                continue
            if total_limit > 0:
                usage_to_limit_ratio = total_usage / total_limit
            total_normalized_usage.append(usage_to_limit_ratio)
        mean = np.mean(total_normalized_usage)
        standard_deviation = statistics.stdev(total_normalized_usage)
        predicted_peak = min(1.0, mean + self.n * standard_deviation)
        current_total_limit = 0
        for vm_state_and_num_sample in vm_states_and_num_samples:
            if len(vm_state_and_num_sample.vm_state.limit) > 0:
                current_total_limit += vm_state_and_num_sample.vm_state.limit[0]
        return predicted_peak * current_total_limit class LimitPredictor(StatefulPredictor):
    def __init__(self, config):
        super().__init__(config)

    def CreateState(self, vm_info):
        return _State()

    def UpdateState(self, vm_measure, vm_state):
        vm_state.limit = vm_measure["sample"]["abstract_metrics"]["limit"]

    def Predict(self, vm_states_and_num_samples):

        vms_limits = []
        for vm_state_and_num_sample in vm_states_and_num_samples:
            vms_limits.append(vm_state_and_num_sample.vm_state.limit)

        predicted_peak = sum(vms_limits)

        return predicted_peak
from past.builtins import unicode
from google.protobuf import text_format 
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.options.pipeline_options import SetupOptions
from apache_beam.io.gcp.internal.clients import bigquery
from simulator.align_by_time import AlignByTime
from simulator.set_abstract_metrics import SetAbstractMetrics
from simulator.filter_vmsample import FilterVMSample
from simulator.reset_and_shift_simulated_time import ResetAndShiftSimulatedTime
from simulator.set_scheduler import SetScheduler
from simulator.fortune_teller import CallFortuneTellerRunner
from simulator.fortune_teller_factory import PredictorFactory
from simulator.config_pb2 import SimulationConfig
from simulator.avg_predictor import AvgPredictor
from simulator.max_predictor import MaxPredictor
from simulator.per_vm_percentile_predictor import PerVMPercentilePredictor
from simulator.per_machine_percentile_predictor import PerMachinePercentilePredictor
from simulator.n_sigma_predictor import NSigmaPredictor
from simulator.limit_predictor import LimitPredictor
import argparse
import apache_beam as beam
import json
import logging
        )
        samples_with_abstract_metrics = (
            pipeline
            | "Query Samples With Abstract Metrics Table"
            >> beam.io.Read(
                beam.io.BigQuerySource(
                    query=samples_with_abstract_metrics_query, use_standard_sql=True
                )
            )
        )

    # Resetting and Shifting
    if configs.reset_and_shift.reset_time_to_zero == True:
        samples_with_reset_and_shift = ResetAndShiftSimulatedTime(
            samples_with_abstract_metrics, configs
        )
    else:
        samples_with_reset_and_shift = samples_with_abstract_metrics

    if configs.samples_with_reset_and_shift.HasField("input"):
        pipeline = beam.Pipeline(options=pipeline_options)
        samples_with_reset_and_shift_query = "SELECT * FROM {}.{}".format(
            configs.samples_with_reset_and_shift.input.dataset,
            configs.samples_with_reset_and_shift.input.table,
        )
        samples_with_reset_and_shift = (
            pipeline
            | "Query Samples With Time Reset and Shift"
            >> beam.io.Read(
                beam.io.BigQuerySource(
                    query=samples_with_reset_and_shift_query, use_standard_sql=True
                )
            )
        )

    # Setting Scheduler
    scheduled_samples = SetScheduler(samples_with_reset_and_shift, configs)

    if configs.scheduled_samples.HasField("input"):
        pipeline = beam.Pipeline(options=pipeline_options)
        scheduled_samples_query = "SELECT * FROM {}.{}".format(
            configs.scheduled_samples.input.dataset,
            configs.scheduled_samples.input.table,
        )
        scheduled_samples = pipeline | "Query Scheduled Samples" >> beam.io.Read(
            beam.io.BigQuerySource(query=scheduled_samples_query, use_standard_sql=True)
        )

    # Calling FortuneTeller Runner
    CallFortuneTellerRunner(scheduled_samples, configs)

    # Saving Filtered Samples
    schema_vmsample_file = open("simulator/schema_vmsample.json")
    schema_vmsample = json.load(schema_vmsample_file)
    if configs.filtered_samples.HasField("output"):
        filtered_samples | "Write Filtered Samples" >> beam.io.WriteToBigQuery(
            "{}.{}".format(
                configs.filtered_samples.output.dataset,
                configs.filtered_samples.output.table,
            ),
            schema=schema_vmsample,
            write_disposition=beam.io.BigQueryDisposition.WRITE_TRUNCATE,
            create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
        )

    # Saving Time Aligned Samples
    schema_simulated_vmsample_file = open("simulator/schema_simulated_vmsample.json")
    schema_simulated_vmsample = json.load(schema_simulated_vmsample_file)
    if configs.time_aligned_samples.HasField("output"):
        time_aligned_samples | "Write Time Aligned Samples" >> beam.io.WriteToBigQuery(
            "{}.{}".format(
                configs.time_aligned_samples.output.dataset,
                configs.time_aligned_samples.output.table,
            ),
            schema=schema_simulated_vmsample,
            write_disposition=beam.io.BigQueryDisposition.WRITE_TRUNCATE,
            create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
        )

    # Saving Samples with Abstract Metrics
    if configs.samples_with_abstract_metrics.HasField("output"):
        samples_with_abstract_metrics | "Write Samples With Abstract Metrics" >> beam.io.WriteToBigQuery(
            "{}.{}".format(
                configs.samples_with_abstract_metrics.output.dataset,
                configs.samples_with_abstract_metrics.output.table,
            ),
            schema=schema_simulated_vmsample,
            write_disposition=beam.io.BigQueryDisposition.WRITE_TRUNCATE,
            create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
        )

    # Saving Samples with Time Reset and Shift
    if configs.samples_with_reset_and_shift.HasField("output"):
        samples_with_reset_and_shift | "Write Reset and Shifted Samples" >> beam.io.WriteToBigQuery(
            "{}.{}".format(
                configs.samples_with_reset_and_shift.output.dataset,
                configs.samples_with_reset_and_shift.output.table,
            ),
            schema=schema_simulated_vmsample,
            write_disposition=beam.io.BigQueryDisposition.WRITE_TRUNCATE,
            create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
        )

    # Saving Scheduled Samples
    if configs.scheduled_samples.HasField("output"):
        scheduled_samples | "Write Scheduled Samples" >> beam.io.WriteToBigQuery(
            "{}.{}".format(
                configs.scheduled_samples.output.dataset,
                configs.scheduled_samples.output.table,
            ),
            schema=schema_simulated_vmsample,
            write_disposition=beam.io.BigQueryDisposition.WRITE_TRUNCATE,
            create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
        )

    result = pipeline.run()
    result.wait_until_finish()


if __name__ == "__main__":

    PredictorFactory().RegisterPredictor(