Minimal test version

app.py

@@ -1,301 +1,101 @@
 """
 HF Spaces app for VLIW kernel optimization via RL.
-Deploy to HF Spaces Pro (A10G GPU).
+Minimal working version.
 """
-import os
-import sys
-import re
-import threading
-import time
-from datetime import datetime
-
 import gradio as gr
 
-# Thread lock for safe state access
-training_state_lock = threading.Lock()
-
-# Constants
-BASELINE_CYCLES = 147734
-TARGET_CYCLES = 1363
-
-# Training state
-training_state = {
-    "running": False,
-    "best_cycles": BASELINE_CYCLES,
-    "best_code": None,
-    "log": [],
-    "start_time": None,
-}
-
-SYSTEM_PROMPT = '''Write optimized VLIW/SIMD kernel. OUTPUT ONLY ONE ```python CODE BLOCK.
-
-ARCHITECTURE: 12 ALU + 6 VALU (VLEN=8) + 2 load + 2 store + 1 flow slots per cycle.
-
-API:
-- alloc_scratch(name, length) -> addr
-- add(engine, slot): engine in {alu, valu, load, store, flow}
-- valu ops work on 8 elements at once
-- build(slots, vliw=True): pack into VLIW bundle
-
-ALGORITHM: 16 rounds x 256 items, hash tree traversal.
-
-OPTIMIZATION: Use vload/vstore (8 elements), pack 6 VALU ops/cycle, unroll loops.
-
-Write complete code with OptimizedKernelBuilder class and run() function.
-'''
-
-
-def extract_code_block(text: str) -> str:
-    """Extract python code from markdown."""
-    pattern = r"```python\s*(.*?)```"
-    matches = re.findall(pattern, text, re.DOTALL)
-    if matches:
-        return matches[-1].strip()
-    pattern = r"```\s*(.*?)```"
-    matches = re.findall(pattern, text, re.DOTALL)
-    if matches:
-        return matches[-1].strip()
-    return text.strip()
+# Check all imports at startup
+startup_log = []
 
-
-def simple_verify(code: str) -> dict:
-    """Simple verification without simulator."""
-    if not code:
-        return {"score": 0.0, "correctness": 0.0, "msg": "Empty"}
-    if "def run" not in code:
-        return {"score": 0.0, "correctness": 0.0, "msg": "No run()"}
-    if "OptimizedKernelBuilder" not in code:
-        return {"score": 0.0, "correctness": 0.0, "msg": "No class"}
-    if "build_kernel" not in code:
-        return {"score": 0.0, "correctness": 0.0, "msg": "No build_kernel"}
-    if "self.add" not in code:
-        return {"score": 0.1, "correctness": 0.5, "msg": "Structural OK"}
-    # Bonus for using vector ops
-    score = 0.2
-    if "vload" in code or "vstore" in code:
-        score += 0.3
-    if "valu" in code:
-        score += 0.3
-    return {"score": score, "correctness": 1.0, "msg": "Good structure"}
-
-
-def log(msg: str):
-    """Thread-safe logging."""
-    timestamp = datetime.now().strftime("%H:%M:%S")
-    formatted = f"[{timestamp}] {msg}"
-    with training_state_lock:
-        training_state["log"].append(formatted)
-    print(formatted)
-
-
-def reward_function(completions: list[str], **kwargs) -> list[float]:
-    """Compute rewards."""
-    rewards = []
-    for completion in completions:
-        try:
-            code = extract_code_block(completion)
-            result = simple_verify(code)
-            reward = result["score"]
-            if result["correctness"] > 0.5:
-                reward += 0.5
-                with training_state_lock:
-                    if not training_state["best_code"] or len(code) > len(training_state["best_code"] or ""):
-                        training_state["best_code"] = code
-                        log(f"New best code (score: {reward:.2f})")
-            rewards.append(reward)
-        except Exception as e:
-            rewards.append(0.0)
-    return rewards
+def check_import(name, import_fn):
+    try:
+        result = import_fn()
+        startup_log.append(f"✓ {name}: {result}")
+        return True
+    except Exception as e:
+        startup_log.append(f"✗ {name}: {str(e)[:80]}")
+        return False
+
+check_import("torch", lambda: __import__("torch").__version__)
+check_import("transformers", lambda: __import__("transformers").__version__)
+check_import("datasets", lambda: __import__("datasets").__version__)
+check_import("peft", lambda: __import__("peft").__version__)
+check_import("trl", lambda: __import__("trl").__version__)
+
+try:
+    from trl import GRPOConfig, GRPOTrainer
+    startup_log.append("✓ GRPOTrainer: OK")
+except Exception as e:
+    startup_log.append(f"✗ GRPOTrainer: {e}")
+
+try:
+    import torch
+    if torch.cuda.is_available():
+        startup_log.append(f"✓ CUDA: {torch.cuda.get_device_name(0)}")
+    else:
+        startup_log.append("✗ CUDA: Not available")
+except Exception as e:
+    startup_log.append(f"✗ CUDA check: {e}")
 
 
-def run_training(model_name: str, num_steps: int, batch_size: int, lr: float, lora_rank: int):
-    """Main training loop."""
-    with training_state_lock:
-        training_state["running"] = True
-        training_state["best_cycles"] = BASELINE_CYCLES
-        training_state["best_code"] = None
-        training_state["log"] = []
-        training_state["start_time"] = time.time()
+def get_status():
+    return "\n".join(startup_log)
 
-    log(f"Starting: {model_name}")
-    log(f"Steps: {num_steps}, Batch: {batch_size}, LR: {lr}")
 
+def test_train(model_name, steps):
+    """Test training function."""
     try:
         import torch
         from datasets import Dataset
-        from transformers import AutoTokenizer, BitsAndBytesConfig, TrainerCallback
+        from transformers import AutoTokenizer
         from peft import LoraConfig
         from trl import GRPOConfig, GRPOTrainer
 
-        if torch.cuda.is_available():
-            log(f"GPU: {torch.cuda.get_device_name(0)}")
-        else:
-            log("WARNING: No GPU!")
+        log = [f"Testing with {model_name}, {steps} steps"]
+        log.append(f"CUDA available: {torch.cuda.is_available()}")
 
-        log("Loading tokenizer...")
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
+        log.append("Tokenizer loaded")
 
-        prompt = f"{SYSTEM_PROMPT}\n\nCURRENT: {BASELINE_CYCLES} cycles. TARGET: <{TARGET_CYCLES}."
-        dataset = Dataset.from_dict({"prompt": [prompt] * 16})
-
-        peft_config = LoraConfig(
-            r=lora_rank,
-            lora_alpha=lora_rank * 2,
-            lora_dropout=0.05,
-            bias="none",
-            task_type="CAUSAL_LM",
-            target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
-        )
-
-        output_dir = f"./output/{datetime.now().strftime('%Y%m%d-%H%M%S')}"
-        os.makedirs(output_dir, exist_ok=True)
+        dataset = Dataset.from_dict({"prompt": ["Write hello world"] * 4})
+        log.append("Dataset created")
 
-        training_args = GRPOConfig(
-            output_dir=output_dir,
-            num_train_epochs=num_steps,
-            per_device_train_batch_size=batch_size,
-            learning_rate=lr,
-            logging_steps=1,
-            save_steps=max(1, num_steps // 5),
-            max_completion_length=1024,
-            temperature=0.7,
-            num_generations=2,
-            beta=0.1,
-            bf16=True,
+        # Just check config creation
+        config = GRPOConfig(
+            output_dir="./test",
+            num_train_epochs=1,
+            per_device_train_batch_size=1,
             report_to="none",
         )
+        log.append("Config created")
 
-        quant_config = None
-        if "7B" in model_name or "7b" in model_name:
-            log("Using 4-bit quantization")
-            quant_config = BitsAndBytesConfig(
-                load_in_4bit=True,
-                bnb_4bit_compute_dtype=torch.bfloat16,
-                bnb_4bit_use_double_quant=True,
-                bnb_4bit_quant_type="nf4",
-            )
-
-        log("Loading model...")
-        model_kwargs = {}
-        if quant_config:
-            model_kwargs["quantization_config"] = quant_config
-
-        class StopCallback(TrainerCallback):
-            def on_step_end(self, args, state, control, **kwargs):
-                if not training_state["running"]:
-                    log("Stopping...")
-                    control.should_training_stop = True
-                return control
-
-        trainer = GRPOTrainer(
-            model=model_name,
-            reward_funcs=[reward_function],
-            args=training_args,
-            train_dataset=dataset,
-            peft_config=peft_config,
-            processing_class=tokenizer,
-            model_init_kwargs=model_kwargs,
-            callbacks=[StopCallback()],
-        )
-
-        log("Model loaded! Training...")
-        trainer.train()
-        log("Training complete!")
-
-        trainer.save_model(os.path.join(output_dir, "final"))
-        log(f"Saved to {output_dir}")
-
-        if training_state["best_code"]:
-            with open(os.path.join(output_dir, "best_code.py"), "w") as f:
-                f.write(training_state["best_code"])
+        return "\n".join(log)
 
     except Exception as e:
         import traceback
-        log(f"ERROR: {e}")
-        log(traceback.format_exc()[:500])
-
-    finally:
-        with training_state_lock:
-            training_state["running"] = False
-            elapsed = time.time() - (training_state["start_time"] or time.time())
-        log(f"Time: {elapsed/60:.1f} min")
-
-
-def start_training(model_name, num_steps, batch_size, lr, lora_rank):
-    if training_state["running"]:
-        return "Already running!"
-    thread = threading.Thread(
-        target=run_training,
-        args=(model_name, int(num_steps), int(batch_size), float(lr), int(lora_rank)),
-        daemon=False
-    )
-    thread.start()
-    return "Training started!"
-
+        return f"Error: {e}\n\n{traceback.format_exc()}"
 
-def stop_training():
-    with training_state_lock:
-        training_state["running"] = False
-    return "Stop signal sent."
 
-
-def get_status():
-    with training_state_lock:
-        if not training_state["start_time"]:
-            return "### Not started\nClick Start to begin training."
-        elapsed = time.time() - training_state["start_time"]
-        is_running = training_state["running"]
-        logs = training_state["log"][-25:]
-
-    return f"""### {'Running' if is_running else 'Stopped'}
-**Time:** {elapsed/60:.1f} min
-
-**Log:**
-```
-{chr(10).join(logs)}
-```"""
-
-
-def get_best_code():
-    with training_state_lock:
-        return training_state["best_code"] or "# No valid code yet"
-
-
-# UI
-with gr.Blocks(title="VLIW Optimizer") as demo:
-    gr.Markdown("# VLIW Kernel Optimizer via RL")
-    gr.Markdown(f"**Baseline:** {BASELINE_CYCLES:,} | **Target:** {TARGET_CYCLES:,}")
+with gr.Blocks(title="VLIW Test") as demo:
+    gr.Markdown("# VLIW Optimizer - Test Mode")
 
     with gr.Row():
-        with gr.Column():
-            model = gr.Dropdown(
-                ["Qwen/Qwen2.5-Coder-3B-Instruct", "Qwen/Qwen2.5-Coder-1.5B-Instruct"],
-                value="Qwen/Qwen2.5-Coder-3B-Instruct",
-                label="Model"
-            )
-            steps = gr.Slider(1, 50, value=10, step=1, label="Steps")
-            batch = gr.Slider(1, 4, value=2, step=1, label="Batch")
-            lr = gr.Number(value=2e-4, label="LR")
-            lora = gr.Slider(8, 32, value=16, step=8, label="LoRA Rank")
-            with gr.Row():
-                start_btn = gr.Button("Start", variant="primary")
-                stop_btn = gr.Button("Stop", variant="stop")
-
-        with gr.Column():
-            status = gr.Markdown("### Not started\nClick Start to begin training.")
-            refresh = gr.Button("Refresh")
+        status_box = gr.Textbox(label="Startup Status", value=get_status(), lines=15)
 
     with gr.Row():
-        code_out = gr.Code(label="Best Code", language="python", lines=20)
-        code_btn = gr.Button("Show Best Code")
+        model = gr.Dropdown(
+            ["Qwen/Qwen2.5-Coder-1.5B-Instruct", "Qwen/Qwen2.5-Coder-3B-Instruct"],
+            value="Qwen/Qwen2.5-Coder-1.5B-Instruct",
+            label="Model"
+        )
+        steps = gr.Slider(1, 5, value=1, label="Steps")
+
+    test_btn = gr.Button("Test Training Setup", variant="primary")
+    output = gr.Textbox(label="Output", lines=15)
 
-    start_btn.click(start_training, [model, steps, batch, lr, lora], [status])
-    stop_btn.click(stop_training, outputs=[status])
-    refresh.click(get_status, outputs=[status])
-    code_btn.click(get_best_code, outputs=[code_out])
-    demo.load(get_status, outputs=[status], every=5)
+    test_btn.click(test_train, [model, steps], [output])
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)