fix: dervies speedup now

.gitignore

@@ -0,0 +1,3 @@
+__pycache__
+venv/
+.env

app.py

@@ -9,7 +9,20 @@ import traceback
 
 load_dotenv()
 
-def ui(task_id:str, max_steps:int, openai_api_key:str)-> Iterator[Tuple[str,str]]:
+def task_baseline_code(task_id: str) -> str:
+    return TASKS[task_id]["baseline_code"]
+
+def extract_code(text: str) -> str:
+    if "```" not in text:
+        return text
+    start = text.find("```")
+    end = text.rfind("```")
+    chunk = text[start + 3 : end]
+    if chunk.startswith("cuda") or chunk.startswith("cpp"):
+        return chunk.split("\n", 1)[1]
+    return chunk
+
+def ui(task_id:str, kernel_code: str, max_steps:int, openai_api_key:str)-> Iterator[Tuple[str,str]]:
     log= []
     env=KernelOptimization_env()
     api_key = openai_api_key or os.getenv("OPENAI_API_KEY")
@@ -17,9 +30,13 @@ def ui(task_id:str, max_steps:int, openai_api_key:str)-> Iterator[Tuple[str,str]
         yield "ERROR: Missing OPENAI_API_KEY", ""
         return
 
-    model = os.getenv("MODEL_NAME", "gpt-4")
-    client = OpenAI(api_key=api_key, base_url=os.getenv("API_BASE_URL", "https://api.openai.com/v1"))
+    model = os.getenv("MODEL_NAME", "llama-3.2-3b")
+    client = OpenAI(api_key=api_key, base_url=os.getenv("API_BASE_URL", "https://api.oxlo.ai/v1"))
     obs = env.reset(task_id=task_id)["observation"]
+    if kernel_code and kernel_code.strip():
+        custom_code = kernel_code.strip()
+        env.state.best_code = custom_code
+        obs["current_best_code"] = custom_code
     best_code = obs["current_best_code"]
     log.append(f"Task: {obs['task_name']}")
 
@@ -34,7 +51,8 @@ def ui(task_id:str, max_steps:int, openai_api_key:str)-> Iterator[Tuple[str,str]
                     {"role": "user", "content": prompt},
                 ],
             )
-            code = (res.choices[0].message.content or "").strip() or obs["current_best_code"]
+            raw = (res.choices[0].message.content or "").strip()
+            code = extract_code(raw).strip() or obs["current_best_code"]
             step = env.step(Action(optimized_code=code, strategy="ui_proposed"))
             obs = step.observation.model_dump()
             best_code = obs["current_best_code"]
@@ -49,13 +67,20 @@ def ui(task_id:str, max_steps:int, openai_api_key:str)-> Iterator[Tuple[str,str]
 with gr.Blocks(title="CUDA Kernel Optimizer") as demo:
     gr.Markdown("CUDA Kernel Optimizer - OpenEnv-aligned workflow")
     task = gr.Dropdown(choices=list(TASKS.keys()), value="vector_add_easy", label="Task")
+    kernel_input = gr.Code(
+        label="Kernel Code (editable, used as optimization input)",
+        language="cpp",
+        lines=16,
+        value=TASKS["vector_add_easy"]["baseline_code"],
+    )
     steps = gr.Slider(minimum=1, maximum=12, value=6, step=1, label="Max Steps")
     key = gr.Textbox(label="OpenAI API Key", type="password", placeholder="sk-...")
     run = gr.Button("Run Optimization", variant="primary")
     logs = gr.Textbox(label="Logs", lines=14)
     code = gr.Code(label="Best Code", language="cpp", lines=16)
-    run.click(ui, inputs=[task, steps, key], outputs=[logs, code])
+    task.change(task_baseline_code, inputs=[task], outputs=[kernel_input])
+    run.click(ui, inputs=[task, kernel_input, steps, key], outputs=[logs, code])
 
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)
+    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)

env_server.py

@@ -112,8 +112,17 @@ class KernelOptimization_env:
         self.current_task_id =task_id or random.choice(list(TASKS.keys()))
         task= TASKS[self.current_task_id]
         self.state =EnvState(initialized=True, task_id=self.current_task_id, step_count=0, max_steps=task["max_steps"], total_reward=0.0, best_code=task["baseline_code"], best_speedup=1.0, completed_checks=[], action_history=[])
-
-        return {"observation": to_observation(self.current_task_id, self.state).model_dump()}
+        return {
+            "observation": to_observation(self.current_task_id, self.state).model_dump(),
+            "info": {
+                "task_id": self.current_task_id,
+                "task_name": task["name"],
+                "difficulty": task["difficulty"],
+                "max_steps": task["max_steps"],
+                "target_speedup": task["target_speedup"],
+                "checks": task["checks"],
+            },
+        }
     
     def step(self, action:Action) ->StepResult:
         if not self.state.initialized or not self.current_task_id:
@@ -129,7 +138,12 @@ class KernelOptimization_env:
         completed.update(newly_completed)
         self.state.completed_checks = sorted(completed)
 
-        est_speedup = self.current_task_id, completed, compile_ok
+        completion_ratio = len(completed) / max(len(TASKS[self.current_task_id]["checks"]), 1)
+        max_reasonable_speedup = 1.0 + completion_ratio * 3.0
+        if action.expected_speedup is None:
+            est_speedup = round(max_reasonable_speedup, 3)
+        else:
+            est_speedup = round(max(1.0, min(action.expected_speedup, max_reasonable_speedup)), 3)
         if est_speedup > self.state.best_speedup:
             self.state.best_speedup = est_speedup
             self.state.best_code = code
@@ -199,4 +213,4 @@ def step(request: StepRequest):
 
 @app.get("/state")
 def state():
-    return env.state_dict()
+    return env.state_dict()

openenv.yaml

@@ -17,7 +17,7 @@ metadata:
   author: aaloksan
 
 tasks:
-  - id: vector_addition_easy
+  - id: vector_add_easy
     name: "Vector Addition Kernel Optimization"
     difficulty: easy
     objective: "Improve memory throughput while preserving correctness."
@@ -39,7 +39,7 @@ interfaces:
   reset:
     method: POST
     path: /reset
-    returns: initial observation and info
+    returns: initial observation, metadata info
   step:
     method: POST
     path: /step
@@ -52,4 +52,4 @@ interfaces:
 baseline:
   script: inference.py
   model_env_var: MODEL_NAME
-  api_key_env_var: OPENAI_API_KEY
+  api_key_env_var: OPENAI_API_KEY

openenv_train.py

@@ -11,8 +11,10 @@ class KernelOptTool:
         self.reward = 0.0
         self.done = False
     
-    def reset(self, **kwargs) ->str|None:
-        task_id =kwargs.get("task_id")
+    def reset(self, **kwargs) -> str:
+        task_id = kwargs.get("task_id")
+        if task_id is None and isinstance(kwargs.get("sample"), dict):
+            task_id = kwargs["sample"].get("task_id")
         result = self.env.reset(task_id=task_id)
         obs = result["observation"]
         self.reward = 0.0
@@ -24,10 +26,12 @@ class KernelOptTool:
             "Use tools to submit improved code."
         )
 
-    def submit_optiization(self, optimized_code:str, strategy:str ="")->str:
+    def submit_optimization(self, optimized_code: str, strategy: str = "", expected_speedup: float | None = None) -> str:
         if self.done:
             raise ValueError("Episode is already done.")
-        result = self.env.step(Action(optimized_code=optimized_code, strategy=strategy))
+        result = self.env.step(
+            Action(optimized_code=optimized_code, strategy=strategy, expected_speedup=expected_speedup)
+        )
         self.reward = result.reward.value
         self.done = result.done
         obs = result.observation
@@ -37,8 +41,12 @@ class KernelOptTool:
             f"pending_checks={obs.pending_checks}, done={result.done}"
         )
 
-def reward_func(environmnets, **kwargs)-> List[float]:
-    return [env.reward for env in environmnets]
+    # Backward-compatible alias
+    def submit_optiization(self, optimized_code: str, strategy: str = "") -> str:
+        return self.submit_optimization(optimized_code=optimized_code, strategy=strategy)
+
+def reward_func(environments, **kwargs) -> List[float]:
+    return [env.reward for env in environments]
 
 def build_dataset(repeats_per_task:int=32)-> Dataset:
     prompts, task_ids = [], []
@@ -67,4 +75,4 @@ def main():
 # trainer = GRPOTrainer(model =model_name, train_dataset=dataset, reward_funcs =reward_func, env_factory=KernelOptTool)
 
 if __name__ == "__main__":
-    main()
+    main()

pyproject.toml

@@ -16,6 +16,7 @@ dependencies = [
   "openenv-core==0.1.1",
   "python-dotenv>=1.0.0",
   "pillow>=10.4.0,<11",
+  "audioop-lts>=0.2.2; python_version >= '3.13'",
   "gradio>=4.44.0",
   "datasets>=2.20.0",
   "trl>=0.12.0"

requirements.txt

@@ -10,3 +10,4 @@ pytest>=7.4.0
 pyyaml>=6.0.0
 python-dotenv
 pillow>=10.4.0,<11
+audioop-lts>=0.2.2; python_version >= "3.13"