manager.py

import time
import hashlib
from kubernetes import client, config
import json
import random
import yaml
import os
from history import *
from es_utils import *

from agent import ReinforceAgent


config.load_incluster_config()
apps_v1 = client.AppsV1Api()
core_v1 = client.CoreV1Api()

NAMESPACE = "rmalves"
DEPLOYMENT_NAME = "collector"
CONFIGMAP_NAME = "collector-config"
POLICIES_FILE = "tail_sampling_policies.json"

def generate_config(selected_policies, config_hash):
    config_dict = {
        "receivers": {
            "otlp": {
                "protocols": {
                    "http": {"endpoint": "0.0.0.0:4321"}
                }
            }
        },
        "processors": {
            "tail_sampling": {
                "decision_wait": "40s",
                "num_traces": 15000,
                "expected_new_traces_per_sec": 1000,
                "policies": selected_policies
            },
            
            "attributes": {
                "actions": [
                    {
                        "key": "experiment_hash",
                        "value": config_hash,
                        "action": "insert"
                    }
                ]
            }
        },
        "exporters": {
            "debug": {"verbosity": "detailed"},
            "otlphttp": {"endpoint": "http://jaeger:4318"},
            "prometheus": {"endpoint": "0.0.0.0:9464"}
        },
        "service": {
            "pipelines": {
                "traces": {
                    "receivers": ["otlp"],
                    "processors": ["tail_sampling", "attributes"],
                    "exporters": ["otlphttp"]
                },
                "metrics": {
                    "receivers": ["otlp"],
                    "exporters": ["prometheus"]
                }
            }
        }
    }

    return yaml.dump(config_dict)

def update_configmap(config_yaml):
    cm_body = client.V1ConfigMap(
        metadata=client.V1ObjectMeta(name=CONFIGMAP_NAME, namespace=NAMESPACE),
        data={"config.yaml": config_yaml}
    )
    try:
        core_v1.replace_namespaced_config_map(CONFIGMAP_NAME, NAMESPACE, cm_body)

    except client.exceptions.ApiException as e:
        if e.status == 404:
            core_v1.create_namespaced_config_map(NAMESPACE, cm_body)
        else:
            raise
def rolling_update_deployment(config_yaml, config_hash):
    patch = {
        "spec": {
            "template": {
                "metadata": {
                    "annotations": {
                        "config-hash": config_hash
                    }
                }
            }
        }
    }
    apps_v1.patch_namespaced_deployment(
        name=DEPLOYMENT_NAME, namespace=NAMESPACE, body=patch
    )

def wait_for_rollout_ready():
    while True:
        deployment = apps_v1.read_namespaced_deployment(DEPLOYMENT_NAME, NAMESPACE)
        desired = deployment.spec.replicas
        available = deployment.status.available_replicas or 0 
        if available >= desired:
            return
        time.sleep(2)
def trace_penalty_function(traces, C, k=25, midpoint=0.20):
    x = traces / C
    return 1 / (1 + math.exp(-k * (x - midpoint)))

def reward_function(entropy, traces, alpha=1.0, beta=1.0, C = 10000, lambd=3.0):
    norm_entropy = entropy/10

    trace_penalty = trace_penalty_function(traces, C)

    return alpha * norm_entropy - beta * trace_penalty

if __name__ == "__main__":

    MAX_NUMBER_EPISODES = 300
    current_episode = 0

    current_test = 0
    MAX_NUMBER_OF_TESTS = 1
    history_buffer = []

    with open(POLICIES_FILE, "r") as f:
        all_policies = json.load(f)

    agent = ReinforceAgent(num_policies = len(all_policies))
    
    first = True

    while True:
        current_episode = current_episode + 1
        if first:
            old_hash = "jausj"
            first = False
        else:
            old_hash = config_hash
    
        selected_policies, selected_actions = agent.select_actions(all_policies)
        
        policies_str = json.dumps(selected_policies, sort_keys = True)
        timestamp = str(time.time())
        hash_input = policies_str + timestamp
        config_hash = hashlib.sha256(hash_input.encode()).hexdigest()[:8] 
        config_yaml = generate_config(selected_policies, config_hash)

        update_configmap(config_yaml)
        rolling_update_deployment(config_yaml, config_hash)

        wait_for_rollout_ready()

        entropia, number_of_traces = export_traces_by_hash(old_hash)

        reward = reward_function(entropia, number_of_traces)
        agent.update(selected_policies, reward, selected_actions)
        print(f"Hash: {config_hash}, reward: {reward}, Entropia: {entropia}, Número de traces: {number_of_traces}")

        history_buffer.append({
            "episode": current_episode,
            "hash": old_hash,
            "entropy": entropia,
            "reward": reward,
            "number_of_traces": number_of_traces,
        })

        if current_episode < MAX_NUMBER_EPISODES:
            time.sleep(60)
        else:
            with open("episodes_history_" + str(current_test) + ".json", "w") as f:
                json.dump(history_buffer, f, indent=2)
            agent.save_policies(current_test)
            first = True
            agent = ReinforceAgent(num_policies = len(all_policies))
            current_test = current_test + 1
            history_buffer = []
            current_episode = 0
            if current_test >= MAX_NUMBER_OF_TESTS:
                while True:
                    time.sleep(100000)