Debug version to check imports

app.py

@@ -1,548 +1,100 @@
 """
 HF Spaces app for VLIW kernel optimization via RL.
-Deploy to HF Spaces Pro (A10G GPU).
-
-This is self-contained - includes verification logic inline.
+Minimal version for debugging.
 """
 import os
 import sys
-import re
-import threading
-import time
-import random
-from datetime import datetime
-
 import gradio as gr
 
-# Thread lock for safe state access
-training_state_lock = threading.Lock()
-
-# Add simulator path
-SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
-PERF_TAKEHOME_PATH = os.path.join(SCRIPT_DIR, "original_performance_takehome")
-if os.path.exists(PERF_TAKEHOME_PATH):
-    sys.path.insert(0, PERF_TAKEHOME_PATH)
-
-# Constants
-BASELINE_CYCLES = 147734
-TARGET_CYCLES = 1363
-SCORE_SCALE = 3000.0
-
-# Training state (global)
-training_state = {
-    "running": False,
-    "step": 0,
-    "total_steps": 0,
-    "best_cycles": BASELINE_CYCLES,
-    "best_code": None,
-    "log": [],
-    "start_time": None,
-    "results": [],
-}
-
-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. 1536-word scratch.
-
-API:
-- alloc_scratch(name, length) -> addr
-- scratch_const(val, name) -> addr
-- add(engine, slot): engine in {alu, valu, load, store, flow}
-  - alu: (op, dst, src1, src2) where op in {+,-,*,/,%,^,&,|,==,!=,<,>,<=,>=}
-  - valu: same ops but on vectors (VLEN=8)
-  - load: (load,dst,addr), (vload,dst,addr), (const,dst,val), (vbroadcast,dst,scalar_addr)
-  - store: (store,addr,src), (vstore,addr,src)
-  - flow: (select,dst,cond,t,f), (jump,label), (jump_if_zero,cond,label), (halt,)
-- label(name): mark code position
-- build(slots, vliw=True): pack slots into VLIW bundle
-
-MEMORY: mem[4]=forest_values, mem[5]=inp_indices, mem[6]=inp_values (256 elements each)
-
-ALGORITHM: 16 rounds x 256 items: load idx,val; val=hash(val^tree[idx]); idx=2*idx+(1 or 2 based on val%2); store. Hash is 16 stages using HASH_STAGES constant.
-
-OPTIMIZATION:
-1. Use vload/vstore: process 8 elements per instruction (256/8 = 32 vector iterations)
-2. Pack ops: 6 VALU slots = 6 vector ops per cycle
-3. Unroll: minimize loop overhead
-4. Pipeline: overlap loads with compute
-
-You MUST override build_kernel() with actual instructions. Do NOT just call super().
-'''
-
-
-def extract_code_block(text: str) -> str:
-    """Extract python code from markdown code blocks."""
-    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 imports on startup
+startup_log = []
 
-
-def verify_perf_takehome(generation: str, score_scale: float = SCORE_SCALE) -> dict:
-    """
-    Verify kernel code and return score.
-    Self-contained verification using the simulator.
-    """
+def check_import(name, import_fn):
     try:
-        code = generation.strip()
-
-        if not code:
-            return {"score": 0.0, "correctness": 0.0, "performance": -1000000,
-                    "msg": "Empty code", "cycles": None}
-
-        if "def run" not in code:
-            return {"score": 0.0, "correctness": 0.0, "performance": -1000000,
-                    "msg": "No 'run' function defined", "cycles": None}
-
-        # Build execution environment
-        exec_globals = {
-            "FOREST_HEIGHT": 10,
-            "ROUNDS": 16,
-            "BATCH_SIZE": 256,
-        }
-
-        # Setup imports
-        setup_code = f'''
-import sys
-sys.path.insert(0, "{PERF_TAKEHOME_PATH}")
-from problem import Machine, Tree, Input, build_mem_image, N_CORES, VLEN, reference_kernel2
-from perf_takehome import KernelBuilder, HASH_STAGES, BASELINE
-import random
-'''
-        full_code = setup_code + "\n" + code
-        exec(full_code, exec_globals)
-
-        if "run" not in exec_globals:
-            return {"score": 0.0, "correctness": 0.0, "performance": -1000000,
-                    "msg": "No 'run' function after exec", "cycles": None}
-
-        # Require OptimizedKernelBuilder
-        if "OptimizedKernelBuilder" not in exec_globals:
-            return {"score": 0.0, "correctness": 0.0, "performance": -1000000,
-                    "msg": "No OptimizedKernelBuilder class", "cycles": None}
-
-        # Run verification
-        random.seed(123)
-        from problem import Tree, Input, Machine, build_mem_image, N_CORES, reference_kernel2
-
-        forest = Tree.generate(10)
-        inp = Input.generate(forest, 256, 16)
-        mem = build_mem_image(forest, inp)
-
-        # Get reference output
-        ref_mem = None
-        for ref_mem in reference_kernel2(list(mem)):
-            pass
-
-        if ref_mem is None:
-            return {"score": 0.0, "correctness": 0.0, "performance": -1000000,
-                    "msg": "Reference kernel failed", "cycles": None}
-
-        # Run submitted kernel
-        kb = exec_globals["OptimizedKernelBuilder"]()
-        kb.build_kernel(10, len(forest.values), 256, 16)
-        machine = Machine(list(mem), kb.instrs, kb.debug_info(), n_cores=N_CORES)
-        machine.enable_pause = False
-        machine.enable_debug = False
-        machine.run()
-
-        cycles = machine.cycle
-
-        # Validate cycles
-        if cycles <= 100:
-            return {"score": 0.0, "correctness": 0.0, "performance": -1000000,
-                    "msg": f"Suspiciously low cycles ({cycles})", "cycles": cycles}
-
-        if cycles > 200000:
-            return {"score": 0.0, "correctness": 0.0, "performance": -1000000,
-                    "msg": f"Cycles too high: {cycles}", "cycles": cycles}
-
-        # Compare outputs
-        inp_values_p = ref_mem[6]
-        expected = ref_mem[inp_values_p : inp_values_p + len(inp.values)]
-        actual = machine.mem[inp_values_p : inp_values_p + len(inp.values)]
-
-        if expected != actual:
-            return {"score": 0.0, "correctness": 0.0, "performance": -1000000,
-                    "msg": f"Incorrect output (cycles={cycles})", "cycles": cycles}
-
-        # Success!
-        score = score_scale / cycles
-        return {
-            "score": score,
-            "correctness": 1.0,
-            "performance": -cycles,
-            "msg": f"Success: {cycles} cycles",
-            "cycles": cycles,
-        }
-
+        result = import_fn()
+        startup_log.append(f"✓ {name}: {result}")
+        return True
     except Exception as e:
-        import traceback
-        tb = traceback.format_exc()
-        error_line = tb.strip().split('\n')[-1][:200]
-        return {"score": 0.0, "correctness": 0.0, "performance": -1000000,
-                "msg": f"Error: {error_line}", "cycles": None}
-
-
-def log(msg: str):
-    """Add to training log (thread-safe)."""
-    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 for completions."""
-    rewards = []
-    for completion in completions:
-        try:
-            code = extract_code_block(completion)
-            result = verify_perf_takehome(code)
-            reward = result["score"]
-
-            if result["correctness"] > 0:
-                reward += 1.0
-                cycles = result.get("cycles")
-                if cycles:
-                    with training_state_lock:
-                        training_state["results"].append({
-                            "step": training_state["step"],
-                            "cycles": cycles,
-                            "time": time.time() - (training_state["start_time"] or time.time())
-                        })
-                        if cycles < training_state["best_cycles"]:
-                            training_state["best_cycles"] = cycles
-                            training_state["best_code"] = code
-                            speedup = BASELINE_CYCLES / cycles
-                            log(f"NEW BEST: {cycles:,} cycles ({speedup:.2f}x speedup)")
-
-            rewards.append(reward)
-
-        except Exception as e:
-            log(f"Reward error: {str(e)[:100]}")
-            rewards.append(0.0)
-
-    return rewards
-
-
-def build_prompt(current_cycles: int = BASELINE_CYCLES, last_code: str = "") -> str:
-    """Build training prompt."""
-    prompt = f"""{SYSTEM_PROMPT}
-
-CURRENT: {current_cycles:,} cycles. TARGET: <{TARGET_CYCLES:,} cycles (need {current_cycles//TARGET_CYCLES}x speedup).
-"""
-    if last_code:
-        prompt += f"""
-Previous best attempt:
-```python
-{last_code[:2000]}
-```
-
-Improve this code to reduce cycles further.
-"""
-    else:
-        prompt += """
-Write a complete solution with:
-1. A run() function that returns (cycles, code_string)
-2. An OptimizedKernelBuilder class with build_kernel() method
-"""
-    return prompt
-
-
-def run_training(model_name: str, num_steps: int, batch_size: int, lr: float, lora_rank: int):
-    """Main training loop."""
-    global training_state
-
-    with training_state_lock:
-        training_state["running"] = True
-        training_state["step"] = 0
-        training_state["total_steps"] = num_steps
-        training_state["best_cycles"] = BASELINE_CYCLES
-        training_state["best_code"] = None
-        training_state["log"] = []
-        training_state["results"] = []
-        training_state["start_time"] = time.time()
-
-    log(f"Starting training: {model_name}")
-    log(f"Steps: {num_steps}, Batch: {batch_size}, LR: {lr}, LoRA rank: {lora_rank}")
-
+        startup_log.append(f"✗ {name}: {str(e)[:100]}")
+        return False
+
+# Test imports
+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__)
+check_import("accelerate", lambda: __import__("accelerate").__version__)
+
+# Try GRPO import
+try:
+    from trl import GRPOConfig, GRPOTrainer
+    startup_log.append("✓ GRPOTrainer: imported from trl")
+except ImportError as e:
+    startup_log.append(f"✗ GRPOTrainer from trl: {e}")
     try:
-        import torch
-        from datasets import Dataset
-        from transformers import AutoTokenizer, BitsAndBytesConfig, TrainerCallback
-        from peft import LoraConfig
-        from trl import GRPOConfig, GRPOTrainer
-
-        # Check GPU
-        if torch.cuda.is_available():
-            gpu_name = torch.cuda.get_device_name(0)
-            gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9
-            log(f"GPU: {gpu_name} ({gpu_mem:.1f}GB)")
-        else:
-            log("WARNING: No GPU detected!")
-
-        log("Loading tokenizer...")
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
-
-        # Create dataset
-        prompt = build_prompt(BASELINE_CYCLES, "")
-        dataset = Dataset.from_dict({"prompt": [prompt] * 64})
-
-        # LoRA config
-        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",
-                            "gate_proj", "up_proj", "down_proj"],
-        )
-
-        # Training config
-        output_dir = f"./output/{datetime.now().strftime('%Y%m%d-%H%M%S')}"
-        os.makedirs(output_dir, exist_ok=True)
-
-        training_args = GRPOConfig(
-            output_dir=output_dir,
-            num_train_epochs=num_steps,
-            per_device_train_batch_size=batch_size,
-            gradient_accumulation_steps=4,
-            learning_rate=lr,
-            logging_steps=1,
-            save_steps=10,
-            max_completion_length=2048,
-            max_prompt_length=2048,
-            temperature=0.7,
-            num_generations=4,
-            beta=0.1,
-            bf16=True,
-            report_to="none",
-        )
-
-        # Quantization for 7B model on A10G
-        quant_config = None
-        if "7B" in model_name or "7b" in model_name:
-            log("Using 4-bit quantization for 7B model")
-            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 (this may take a few minutes)...")
-
-        model_kwargs = {}
-        if quant_config:
-            model_kwargs["quantization_config"] = quant_config
-
-        # Create stop callback
-        class StopCallback(TrainerCallback):
-            def on_step_end(self, args, state, control, **kwargs):
-                if not training_state["running"]:
-                    log("Stop signal received, halting training...")
-                    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! Starting training...")
-
-        # Train
-        trainer.train()
-
-        log("Training complete!")
-
-        # Save
-        trainer.save_model(os.path.join(output_dir, "final"))
-        log(f"Model saved to {output_dir}/final")
-
-        # Save best code
-        if training_state["best_code"]:
-            with open(os.path.join(output_dir, "best_code.py"), "w") as f:
-                f.write(training_state["best_code"])
-            log("Best code saved!")
-
-    except Exception as e:
-        import traceback
-        log(f"ERROR: {str(e)}")
-        log(traceback.format_exc())
-
-    finally:
-        with training_state_lock:
-            training_state["running"] = False
-            elapsed = time.time() - training_state["start_time"]
-            best = training_state["best_cycles"]
-        log(f"Total time: {elapsed/60:.1f} minutes")
-        log(f"Best result: {best:,} cycles")
-
-
-def start_training(model_name, num_steps, batch_size, lr, lora_rank):
-    """Start training in background."""
-    if training_state["running"]:
-        return "Training already running!"
-
-    thread = threading.Thread(
-        target=run_training,
-        args=(model_name, int(num_steps), int(batch_size), float(lr), int(lora_rank)),
-        daemon=False  # Non-daemon to ensure training completes
-    )
-    thread.start()
-    return "Training started! Monitor progress below."
-
-
-def stop_training():
-    """Signal training to stop."""
-    with training_state_lock:
-        training_state["running"] = False
-    return "Stop signal sent. Training will stop after current step."
-
-
-def get_status():
-    """Get current status as markdown."""
-    if not training_state["start_time"]:
-        return "### Status: Not started\n\nConfigure settings and click Start Training."
-
-    with training_state_lock:
-        elapsed = time.time() - training_state["start_time"]
-        elapsed_str = f"{elapsed/60:.1f} min"
-        best_cycles = max(training_state["best_cycles"], 1)  # Prevent division by zero
-        is_running = training_state["running"]
-        log_lines = training_state["log"][-15:]
-
-    speedup = BASELINE_CYCLES / best_cycles
-    progress_pct = (1 - best_cycles / BASELINE_CYCLES) * 100
-
-    status = f"""### Status: {'Running' if is_running else 'Stopped'}
-
-| Metric | Value |
-|--------|-------|
-| Elapsed | {elapsed_str} |
-| Best Cycles | **{best_cycles:,}** |
-| Speedup | **{speedup:.2f}x** |
-| Progress to Target | {progress_pct:.1f}% |
-| Target | {TARGET_CYCLES:,} cycles |
-
----
-
-### Recent Log
-```
-{chr(10).join(log_lines)}
-```
-"""
-    return status
-
-
-def get_best_code():
-    """Get best code found."""
-    with training_state_lock:
-        best_code = training_state["best_code"]
-    if best_code:
-        return best_code
-    return "# No valid code found yet.\n# Start training to generate optimized kernels."
-
-
-def get_results_chart():
-    """Get results as simple text chart."""
-    with training_state_lock:
-        results = list(training_state["results"][-20:])
-
-    if not results:
-        return "No results yet."
-
-    lines = ["Cycles over time:", ""]
-    for r in results:
-        bar_len = max(1, int(50 * r["cycles"] / BASELINE_CYCLES))
-        bar = "#" * bar_len
-        lines.append(f"{r['cycles']:>7,} | {bar}")
-
-    return "\n".join(lines)
-
+        from trl.trainer.grpo_trainer import GRPOConfig, GRPOTrainer
+        startup_log.append("✓ GRPOTrainer: imported from trl.trainer.grpo_trainer")
+    except ImportError as e2:
+        startup_log.append(f"✗ GRPOTrainer alt: {e2}")
+
+# Check CUDA
+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:
+    startup_log.append("✗ CUDA: Could not check")
+
+# Check simulator
+try:
+    SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+    PERF_PATH = os.path.join(SCRIPT_DIR, "original_performance_takehome")
+    if os.path.exists(PERF_PATH):
+        sys.path.insert(0, PERF_PATH)
+        from problem import Machine, Tree
+        startup_log.append("✓ Simulator: loaded")
+    else:
+        startup_log.append(f"✗ Simulator: path not found ({PERF_PATH})")
+except Exception as e:
+    startup_log.append(f"✗ Simulator: {e}")
 
-# Build Gradio UI
-with gr.Blocks(title="VLIW Kernel Optimizer", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("""
-    # VLIW Kernel Optimization via Reinforcement Learning
 
-    Train a language model to generate optimized VLIW/SIMD kernels.
+def get_startup_log():
+    return "\n".join(startup_log)
 
-    | Baseline | Target | Goal |
-    |----------|--------|------|
-    | 147,734 cycles | 1,363 cycles | 108x speedup |
-    """)
 
-    with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown("### Configuration")
-
-            model_dropdown = gr.Dropdown(
-                choices=[
-                    "Qwen/Qwen2.5-Coder-7B-Instruct",
-                    "Qwen/Qwen2.5-Coder-3B-Instruct",
-                    "Qwen/Qwen2.5-Coder-1.5B-Instruct",
-                    "deepseek-ai/deepseek-coder-6.7b-instruct",
-                    "codellama/CodeLlama-7b-Instruct-hf",
-                ],
-                value="Qwen/Qwen2.5-Coder-7B-Instruct",
-                label="Model"
-            )
+def dummy_train(model, steps):
+    return f"Would train {model} for {steps} steps\n\nImport status:\n" + get_startup_log()
 
-            steps_slider = gr.Slider(1, 100, value=50, step=1, label="Training Steps")
-            batch_slider = gr.Slider(1, 8, value=4, step=1, label="Batch Size")
-            lr_input = gr.Number(value=2e-4, label="Learning Rate")
-            lora_slider = gr.Slider(8, 64, value=32, step=8, label="LoRA Rank")
 
-            with gr.Row():
-                start_btn = gr.Button("Start Training", variant="primary", size="lg")
-                stop_btn = gr.Button("Stop", variant="stop")
-
-        with gr.Column(scale=2):
-            status_md = gr.Markdown("### Status: Not started")
-            refresh_btn = gr.Button("Refresh", size="sm")
+# Simple UI
+with gr.Blocks(title="VLIW Optimizer") as demo:
+    gr.Markdown("# VLIW Kernel Optimizer - Debug Mode")
+    gr.Markdown("Checking if all imports work...")
 
     with gr.Row():
         with gr.Column():
-            gr.Markdown("### Best Code Found")
-            code_output = gr.Code(language="python", lines=25)
-            code_btn = gr.Button("Show Best Code")
+            status = gr.Textbox(
+                label="Startup Log",
+                value=get_startup_log(),
+                lines=20
+            )
+            refresh_btn = gr.Button("Refresh Status")
+            refresh_btn.click(get_startup_log, outputs=[status])
 
         with gr.Column():
-            gr.Markdown("### Results")
-            results_output = gr.Textbox(lines=15, label="Cycles Progress")
-            results_btn = gr.Button("Show Results")
-
-    # Event handlers
-    start_btn.click(
-        start_training,
-        inputs=[model_dropdown, steps_slider, batch_slider, lr_input, lora_slider],
-        outputs=[status_md]
-    )
-    stop_btn.click(stop_training, outputs=[status_md])
-    refresh_btn.click(get_status, outputs=[status_md])
-    code_btn.click(get_best_code, outputs=[code_output])
-    results_btn.click(get_results_chart, outputs=[results_output])
-
-    # Auto-refresh
-    demo.load(get_status, outputs=[status_md], every=5)
+            model = gr.Dropdown(
+                choices=["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, 10, value=3, label="Steps")
+            train_btn = gr.Button("Test Train", variant="primary")
+            output = gr.Textbox(label="Output", lines=10)
+            train_btn.click(dummy_train, inputs=[model, steps], outputs=[output])
 
 
 if __name__ == "__main__":