File size: 3,064 Bytes
4c614b7
d9dab8d
4c614b7
 
a87d89e
4c614b7
 
 
d755597
 
4c614b7
d755597
4c614b7
d755597
 
4c614b7
d755597
4c614b7
 
 
 
 
ebad63b
d755597
d9dab8d
d755597
 
 
 
 
 
 
a87d89e
d755597
 
ebad63b
d9dab8d
ebad63b
9a0bf91
ebad63b
 
d755597
 
 
 
 
ebad63b
d755597
 
 
 
 
 
 
ebad63b
d755597
 
 
 
d9dab8d
d755597
 
 
4c614b7
ebad63b
 
 
 
 
 
 
 
 
 
 
4c614b7
d755597
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
import os
import asyncio
from dotenv import load_dotenv
from openai import OpenAI
from environment import CloudCostEnv

load_dotenv()

def log_start(task: str, env: str, model: str) -> None:
    print(f"[START] task={task} env={env} model={model}", flush=True)

def log_step(step: int, action: str, reward: float, done: bool, error=None) -> None:
    error_val = error if error else "null"
    done_val = str(done).lower()
    print(f"[STEP] step={step} action={action} reward={reward:.2f} done={done_val} error={error_val}", flush=True)

def log_end(success: bool, steps: int, score: float, rewards: list) -> None:
    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
    print(f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}", flush=True)

SYSTEM_PROMPT = "You are a Cloud Infrastructure AI. Manage the cluster by choosing actions: 0(Stay), 1(Add Server), 2(Remove), 3(Toggle Spot). Reply with ONLY the integer."

async def get_action(state, client, model_name):
    state_data = state.model_dump_json() if hasattr(state, 'model_dump_json') else str(state)
    try:
        response = client.chat.completions.create(
            model=model_name,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": f"Current State: {state_data}"}
            ],
            max_tokens=10
        )
        action_text = response.choices[0].message.content.strip()
        action_int = int(''.join(filter(str.isdigit, action_text))[0])
        return action_int
    except Exception as e:
        return 0

async def run_task(task_id, difficulty, target_score, client, model_name):
    env = CloudCostEnv(task=difficulty)
    log_start(task=task_id, env="CloudEnv-v1", model=model_name)
    rewards = []
    state = await env.reset()
    try:
        for step in range(1, 21):
            action_int = await get_action(state, client, model_name)
            state, reward, done = await env.step(action_int)
            log_step(step=step, action=str(action_int), reward=reward, done=done)
            rewards.append(reward)
            if done:
                break
        avg_score = sum(rewards) / len(rewards) if rewards else 0
        log_end(
            success=(avg_score >= target_score),
            steps=len(rewards),
            score=avg_score,
            rewards=rewards
        )
    except Exception as e:
        log_step(step=0, action="error", reward=0.0, done=True, error=str(e))
    finally:
        await env.close()

async def main():
    api_base = os.getenv("API_BASE_URL")
    api_key = os.getenv("HF_TOKEN")
    model_name = os.getenv("MODEL_NAME")
    client = OpenAI(base_url=api_base, api_key=api_key)

    # Run all 3 tasks — checker needs at least 3
    await run_task("daily_peaks",   "easy",   0.8,  client, model_name)
    await run_task("weekly_budget", "medium", 0.75, client, model_name)
    await run_task("chaos_mode",    "hard",   0.7,  client, model_name)

if __name__ == "__main__":
    asyncio.run(main())