Spaces:
Sleeping
Sleeping
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()) |