diff --git a/MaxKernel/auto_agent/server_utils/common_models.py b/MaxKernel/auto_agent/server_utils/common_models.py new file mode 100644 index 0000000..5e74550 --- /dev/null +++ b/MaxKernel/auto_agent/server_utils/common_models.py @@ -0,0 +1,50 @@ +import logging +from typing import Optional +from pydantic import BaseModel + +class CodeRequest(BaseModel): + code: str + timeout: Optional[int] = 30 + dependencies: Optional[dict] = None + + +class CodeResponse(BaseModel): + output: str + error: Optional[str] = None + exit_code: int + + +class AutotuneRequest(BaseModel): + code_template: str + search_space: dict[str, list] + timeout: Optional[int] = 300 + total_timeout: Optional[int] = None + dependencies: Optional[dict] = None + + +class GetTpuVersionResponse(BaseModel): + tpu_version: str + + +class GetBackendVersionResponse(BaseModel): + backend_version: str + + +def extract_code(code: str) -> str: + """Extracts code from markdown blocks if present.""" + code_content = code.strip() + if code_content.startswith("```python") and code_content.endswith("```"): + lines = code_content.split("\n") + if lines[0].strip() == "```python": + lines = lines[1:] + if lines[-1].strip() == "```": + lines = lines[:-1] + return "\n".join(lines) + elif code_content.startswith("```") and code_content.endswith("```"): + lines = code_content.split("\n") + if lines[0].strip().startswith("```"): + lines = lines[1:] + if lines[-1].strip() == "```": + lines = lines[:-1] + return "\n".join(lines) + return code_content diff --git a/MaxKernel/auto_agent/server_utils/cpu_server.py b/MaxKernel/auto_agent/server_utils/cpu_server.py index 7b362e5..3f6ca4b 100644 --- a/MaxKernel/auto_agent/server_utils/cpu_server.py +++ b/MaxKernel/auto_agent/server_utils/cpu_server.py @@ -7,10 +7,15 @@ from typing import Optional from fastapi import FastAPI, HTTPException -from pydantic import BaseModel from auto_agent.constants import CPU_SERVER_PORT from auto_agent.tools.analyze_profile import analyze_trace +from auto_agent.server_utils.common_models import ( + CodeRequest, + CodeResponse, + GetBackendVersionResponse, + extract_code, +) logging.basicConfig( level=logging.INFO, @@ -27,22 +32,6 @@ profile_semaphore = asyncio.Semaphore(1) -class CodeRequest(BaseModel): - code: str - timeout: Optional[int] = 30 - dependencies: Optional[dict] = None - - -class CodeResponse(BaseModel): - output: str - error: Optional[str] = None - exit_code: int - - -class GetBackendVersionResponse(BaseModel): - backend_version: str - - def get_cpu_env(): """ Returns environment variables that force JAX to use CPU backend. @@ -55,26 +44,6 @@ def get_cpu_env(): return env -def extract_code(code: str) -> str: - """Extracts code from markdown blocks if present.""" - code_content = code.strip() - if code_content.startswith("```python") and code_content.endswith("```"): - lines = code_content.split("\n") - if lines[0].strip() == "```python": - lines = lines[1:] - if lines[-1].strip() == "```": - lines = lines[:-1] - return "\n".join(lines) - elif code_content.startswith("```") and code_content.endswith("```"): - lines = code_content.split("\n") - if lines[0].strip().startswith("```"): - lines = lines[1:] - if lines[-1].strip() == "```": - lines = lines[:-1] - return "\n".join(lines) - return code_content - - @app.get("/health") async def health_check(): return {"status": "healthy", "backend": "cpu"} @@ -100,17 +69,38 @@ async def compilation_test(request: CodeRequest): with open(file_path, "w") as f: f.write(content) + # Write sitecustomize.py to force Pallas interpret mode on CPU fallback + sitecustomize_path = os.path.join(temp_dir, "sitecustomize.py") + with open(sitecustomize_path, "w") as f: + f.write(""" +import sys +try: + import jax.experimental.pallas as pl + original_pallas_call = pl.pallas_call + def mocked_pallas_call(*args, **kwargs): + kwargs['interpret'] = True + return original_pallas_call(*args, **kwargs) + pl.pallas_call = mocked_pallas_call +except Exception: + pass +""") + with open(temp_file_path, "w") as f: f.write(request.code) - # Execute the code in a subprocess with CPU-only environment + + # Execute the code in a subprocess with CPU-only environment and sitecustomize active + env = get_cpu_env() + current_pythonpath = env.get("PYTHONPATH", "") + env["PYTHONPATH"] = f"{temp_dir}:{current_pythonpath}" if current_pythonpath else temp_dir + process = await asyncio.create_subprocess_exec( sys.executable, temp_file_path, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, cwd=temp_dir, - env=get_cpu_env(), # Force CPU backend + env=env, ) try: @@ -484,11 +474,6 @@ async def profile(request: CodeRequest): await process.wait() except Exception as e: logging.error(f"Failed to kill process: {e}") - # Clean up the temporary directory - # try: - # shutil.rmtree(temp_dir) - # except Exception: - # pass logging.info("Profile analysis finished") diff --git a/MaxKernel/auto_agent/server_utils/mock_tpu_server.py b/MaxKernel/auto_agent/server_utils/mock_tpu_server.py new file mode 100644 index 0000000..2b25340 --- /dev/null +++ b/MaxKernel/auto_agent/server_utils/mock_tpu_server.py @@ -0,0 +1,611 @@ +import asyncio +import itertools +import json +import logging +import os +import re +import shutil +import sys +import tempfile +import time +from typing import Optional + +from fastapi import FastAPI, HTTPException + +from auto_agent.constants import TPU_SERVER_PORT +from auto_agent.server_utils.common_models import ( + CodeRequest, + CodeResponse, + AutotuneRequest, + GetTpuVersionResponse, + extract_code, +) +from evaluation.fake_kernel_evaluator import FakeKernelEvaluator + + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", +) + +app = FastAPI(title="Mock TPU Code Execution Server", version="1.0.0") + +# Semaphore to limit concurrent compilation requests +compilation_semaphore = asyncio.Semaphore(1) +correctness_semaphore = asyncio.Semaphore(1) +performance_semaphore = asyncio.Semaphore(1) +profile_semaphore = asyncio.Semaphore(1) +autotune_semaphore = asyncio.Semaphore(1) + + +@app.get("/health") +async def health_check(): + return {"status": "healthy", "mode": "mock"} + + +@app.post("/compilation_test", response_model=CodeResponse) +async def compilation_test(request: CodeRequest): + """Try to execute kernel safely in a subprocess and return the output.""" + logging.info("Starting compilation test (mock server)") + async with compilation_semaphore: + try: + request.code = extract_code(request.code) + temp_dir = tempfile.mkdtemp() + temp_file_path = os.path.join(temp_dir, "run_code.py") + + if request.dependencies: + for filename, content in request.dependencies.items(): + file_path = os.path.join(temp_dir, filename) + os.makedirs(os.path.dirname(file_path), exist_ok=True) + with open(file_path, "w") as f: + f.write(content) + + # Write sitecustomize.py to force Pallas interpret mode on CPU fallback + sitecustomize_path = os.path.join(temp_dir, "sitecustomize.py") + with open(sitecustomize_path, "w") as f: + f.write(""" +import sys +try: + import jax.experimental.pallas as pl + original_pallas_call = pl.pallas_call + def mocked_pallas_call(*args, **kwargs): + kwargs['interpret'] = True + return original_pallas_call(*args, **kwargs) + pl.pallas_call = mocked_pallas_call +except Exception: + pass +""") + + with open(temp_file_path, "w") as f: + f.write(request.code) + + env = os.environ.copy() + current_pythonpath = env.get("PYTHONPATH", "") + env["PYTHONPATH"] = f"{temp_dir}:{current_pythonpath}" if current_pythonpath else temp_dir + + process = await asyncio.create_subprocess_exec( + sys.executable, + temp_file_path, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=temp_dir, + env=env, + ) + + try: + stdout, stderr = await asyncio.wait_for( + process.communicate(), timeout=request.timeout + ) + + output = stdout.decode("utf-8") if stdout else "" + error = stderr.decode("utf-8") if stderr else None + exit_code = process.returncode + + logging.info( + f"Compilation test completed successfully with exit_code: {exit_code}" + ) + return CodeResponse(output=output, error=error, exit_code=exit_code) + + except asyncio.TimeoutError: + process.kill() + await process.wait() + logging.error(f"Compilation test timed out after {request.timeout}s") + raise HTTPException(status_code=408, detail="Code execution timed out") + + except HTTPException: + raise + except Exception as e: + logging.error(f"Compilation test failed with error: {str(e)}") + raise HTTPException(status_code=500, detail=f"Execution error: {str(e)}") + finally: + if "process" in locals() and process.returncode is None: + try: + process.kill() + await process.wait() + except Exception as e: + logging.error(f"Failed to kill process: {e}") + if "temp_dir" in locals(): + try: + shutil.rmtree(temp_dir) + except Exception: + pass + + +@app.post("/correctness_test", response_model=CodeResponse) +async def correctness_test(request: CodeRequest): + """Test correctness by running local subprocess.""" + logging.info("Starting correctness test (mock server)") + async with correctness_semaphore: + try: + request.code = extract_code(request.code) + temp_dir = tempfile.mkdtemp() + temp_file_path = os.path.join(temp_dir, "run_code.py") + + if request.dependencies: + for filename, content in request.dependencies.items(): + file_path = os.path.join(temp_dir, filename) + os.makedirs(os.path.dirname(file_path), exist_ok=True) + with open(file_path, "w") as f: + f.write(content) + + with open(temp_file_path, "w") as f: + f.write(request.code) + + process = await asyncio.create_subprocess_exec( + sys.executable, + temp_file_path, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=temp_dir, + ) + + try: + stdout, stderr = await asyncio.wait_for( + process.communicate(), timeout=request.timeout + ) + + output = stdout.decode("utf-8") if stdout else "" + error = stderr.decode("utf-8") if stderr else None + exit_code = process.returncode + + logging.info( + f"Correctness test completed successfully with exit_code: {exit_code}" + ) + return CodeResponse(output=output, error=error, exit_code=exit_code) + + except asyncio.TimeoutError: + process.kill() + await process.wait() + logging.error(f"Correctness test timed out after {request.timeout}s") + raise HTTPException(status_code=408, detail="Code execution timed out") + + except HTTPException: + raise + except Exception as e: + logging.error(f"Correctness test failed with error: {str(e)}") + raise HTTPException(status_code=500, detail=f"Execution error: {str(e)}") + finally: + if "process" in locals() and process.returncode is None: + try: + process.kill() + await process.wait() + except Exception as e: + logging.error(f"Failed to kill process: {e}") + if "temp_dir" in locals(): + try: + shutil.rmtree(temp_dir) + except Exception: + pass + + +@app.post("/performance_test", response_model=CodeResponse) +async def performance_test(request: CodeRequest): + """Test performance by running local subprocess.""" + logging.info("Starting performance test (mock server)") + async with performance_semaphore: + try: + request.code = extract_code(request.code) + temp_dir = tempfile.mkdtemp() + temp_file_path = os.path.join(temp_dir, "run_code.py") + + if request.dependencies: + for filename, content in request.dependencies.items(): + file_path = os.path.join(temp_dir, filename) + os.makedirs(os.path.dirname(file_path), exist_ok=True) + with open(file_path, "w") as f: + f.write(content) + + with open(temp_file_path, "w") as f: + f.write(request.code) + + process = await asyncio.create_subprocess_exec( + sys.executable, + temp_file_path, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=temp_dir, + ) + + try: + stdout, stderr = await asyncio.wait_for( + process.communicate(), timeout=request.timeout + ) + + output = stdout.decode("utf-8") if stdout else "" + error = stderr.decode("utf-8") if stderr else None + exit_code = process.returncode + + logging.info( + f"Performance test completed successfully with exit_code: {exit_code}" + ) + return CodeResponse(output=output, error=error, exit_code=exit_code) + + except asyncio.TimeoutError: + process.kill() + await process.wait() + logging.error(f"Performance test timed out after {request.timeout}s") + raise HTTPException(status_code=408, detail="Code execution timed out") + + except HTTPException: + raise + except Exception as e: + logging.error(f"Performance test failed with error: {str(e)}") + raise HTTPException(status_code=500, detail=f"Execution error: {str(e)}") + finally: + if "process" in locals() and process.returncode is None: + try: + process.kill() + await process.wait() + except Exception as e: + logging.error(f"Failed to kill process: {e}") + if "temp_dir" in locals(): + try: + shutil.rmtree(temp_dir) + except Exception: + pass + + +@app.post("/unified_test", response_model=CodeResponse) +async def unified_test(request: CodeRequest): + """Simulate correctness and performance with FakeKernelEvaluator.""" + logging.info("Starting mock unified test") + async with performance_semaphore: + try: + request.code = extract_code(request.code) + temp_dir = tempfile.mkdtemp() + temp_file_path = os.path.join(temp_dir, "run_code.py") + + if request.dependencies: + for filename, content in request.dependencies.items(): + file_path = os.path.join(temp_dir, filename) + os.makedirs(os.path.dirname(file_path), exist_ok=True) + with open(file_path, "w") as f: + f.write(content) + + with open(temp_file_path, "w") as f: + f.write(request.code) + + ref_path = None + opt_path = None + if request.dependencies: + for name in request.dependencies.keys(): + if "optimized" in name: + opt_path = os.path.join(temp_dir, name) + elif name == "kernel.py" or name.endswith(".py"): + ref_path = os.path.join(temp_dir, name) + + if not ref_path and request.dependencies: + for name in request.dependencies.keys(): + if name.endswith(".py") and "optimized" not in name: + ref_path = os.path.join(temp_dir, name) + break + + if not ref_path: + ref_path = temp_file_path + if not opt_path: + opt_path = temp_file_path + + evaluator = FakeKernelEvaluator() + eval_result = evaluator.evaluate( + reference_code_path=ref_path, + optimized_code_path=opt_path, + ) + + output = ( + f"CORRECTNESS: {str(eval_result.numerically_correct).upper()}\n" + f"RESULT_TIME: {eval_result.optimized_time_ms} ms\n" + ) + exit_code = ( + 0 + if eval_result.compiled_successfully and eval_result.numerically_correct + else 1 + ) + error = ( + None + if exit_code == 0 + else "Simulated compilation or correctness check failed." + ) + logging.info( + f"Mock unified test completed with exit_code: {exit_code}" + ) + return CodeResponse(output=output, error=error, exit_code=exit_code) + + except Exception as e: + logging.error(f"Unified test failed with error: {str(e)}") + raise HTTPException(status_code=500, detail=f"Execution error: {str(e)}") + finally: + if "temp_dir" in locals(): + try: + shutil.rmtree(temp_dir) + except Exception: + pass + + +@app.post("/profile", response_model=CodeResponse) +async def profile(request: CodeRequest): + """Simulate profile runs and output mock xplane.pb file.""" + logging.info("Starting mock profile") + async with profile_semaphore: + try: + request.code = extract_code(request.code) + temp_dir = tempfile.mkdtemp() + logging.info(f"temp_dir: {temp_dir}") + + if request.dependencies: + for filename, content in request.dependencies.items(): + file_path = os.path.join(temp_dir, filename) + os.makedirs(os.path.dirname(file_path), exist_ok=True) + with open(file_path, "w") as f: + f.write(content) + + temp_file_path = os.path.join(temp_dir, "profile_code.py") + with open(temp_file_path, "w") as temp_file: + temp_file.write(request.code) + + ref_path = None + opt_path = None + if request.dependencies: + for name in request.dependencies.keys(): + if "optimized" in name: + opt_path = os.path.join(temp_dir, name) + elif name == "kernel.py" or name.endswith(".py"): + ref_path = os.path.join(temp_dir, name) + + if not ref_path and request.dependencies: + for name in request.dependencies.keys(): + if name.endswith(".py") and "optimized" not in name: + ref_path = os.path.join(temp_dir, name) + break + + if not ref_path: + ref_path = temp_file_path + if not opt_path: + opt_path = temp_file_path + + evaluator = FakeKernelEvaluator() + eval_result = evaluator.evaluate( + reference_code_path=ref_path, + optimized_code_path=opt_path, + ) + + metrics_path = os.path.join(temp_dir, "evaluation_metrics.json") + hbm_util = 45.0 + compute_util = 15.0 + if os.path.exists(metrics_path): + try: + with open(metrics_path, "r") as f: + metrics_data = json.load(f) + hbm_util = float(metrics_data.get("estimated_hbm_utilization", 45.0)) + compute_util = float( + metrics_data.get("estimated_compute_utilization", 15.0) + ) + except Exception as e: + logging.warning(f"Failed to read evaluation_metrics.json: {e}") + + total_util = hbm_util + compute_util + ratio = hbm_util / total_util if total_util > 0 else 0.75 + + mock_xplane_path = os.path.join(temp_dir, "mock_profile.xplane.pb") + with open(mock_xplane_path, "w") as f: + f.write("MOCK_TRACE_DATA") + + logging.info( + f"Mock profile completed. Ratio: {ratio}, Path: {mock_xplane_path}" + ) + return CodeResponse( + output=json.dumps({"ratio": ratio, "xplane_path": mock_xplane_path}), + error=None, + exit_code=0, + ) + + except Exception as e: + logging.error(f"Profile failed with error: {str(e)}") + raise HTTPException(status_code=500, detail=f"Execution error: {str(e)}") + finally: + # Keep temp_dir alive for parent profile analysis reads of xplane file + pass + + +@app.post("/autotune", response_model=CodeResponse) +async def autotune(request: AutotuneRequest): + """Simulate autotune sweeps using Gemini or fallbacks.""" + logging.info("Starting mock autotune") + async with performance_semaphore: + temp_dir = None + try: + temp_dir = tempfile.mkdtemp() + if request.dependencies: + for filename, content in request.dependencies.items(): + file_path = os.path.join(temp_dir, filename) + os.makedirs(os.path.dirname(file_path), exist_ok=True) + with open(file_path, "w") as f: + f.write(content) + + keys = list(request.search_space.keys()) + values = list(request.search_space.values()) + combinations = list(itertools.product(*values)) + + best_time = float("inf") + best_cfg = None + best_output = "" + all_results = [] + + configs_list = [] + candidates_formatted = [] + for idx, combo in enumerate(combinations): + cfg = dict(zip(keys, combo)) + configs_list.append(cfg) + cfg_str = ", ".join(f"{k}={v}" for k, v in cfg.items()) + candidates_formatted.append(f"Candidate {idx}: {cfg_str}") + candidates_list_str = "\n".join(candidates_formatted) + + ref_code = "" + if request.dependencies: + for name, content in request.dependencies.items(): + if name == "kernel.py" or name.endswith(".py"): + ref_code = content + break + + prompt = f"""You are an expert TPU performance engineer. +Analyze the baseline JAX kernel and the optimized JAX kernel template. +Estimate the execution latency (in milliseconds) on Cloud TPU v4 for each configuration candidate. + +Baseline Reference Code: +{ref_code} + +Optimized Kernel Code Template: +{request.code_template} + +Given: +- The reference baseline latency is 12.0 ms. +- The hardware target is a Cloud TPU v4. + +Predict metrics for each of the following candidate configurations: +{candidates_list_str} + +For each candidate, predict: +1. `latency_ms`: Execution latency in milliseconds. If it triggers compiler errors, shape mismatches, or JAX/Pallas errors, return 999.0. +2. `status`: "success" or "failed". +3. `error`: A brief string of the compiler error if failed, otherwise empty. + +Respond with a single raw JSON matching this schema: +{{ + "results": [ + {{ + "cfg_index": int, + "latency_ms": float, + "status": "string", + "error": "string" + }}, + ... + ] +}} +Do NOT wrap response in markdown blocks. Return only JSON. +""" + gemini_success = False + try: + from google.genai import Client + from google.genai import types + + client = Client() + response = client.models.generate_content( + model="gemini-2.5-flash", + contents=prompt, + config=types.GenerateContentConfig( + response_mime_type="application/json", + temperature=0.1, + ), + ) + response_text = response.text.strip() + data = json.loads(response_text) + results = data.get("results", []) + for res in results: + idx = res.get("cfg_index") + if idx is not None and 0 <= idx < len(configs_list): + cfg = configs_list[idx] + latency = float(res.get("latency_ms", 999.0)) + status = res.get("status", "success") + err_msg = res.get("error", "") + + if status == "success" and latency < 999.0: + all_results.append( + {"cfg": cfg, "time": latency, "status": "success"} + ) + if latency < best_time: + best_time = latency + best_cfg = cfg + best_output = f"CORRECTNESS: TRUE\nRESULT_TIME: {latency} ms" + else: + all_results.append( + { + "cfg": cfg, + "status": "failed", + "error": err_msg or "Simulated compiler error", + "exit_code": 1, + } + ) + gemini_success = len(all_results) > 0 + except Exception as e: + logging.warning( + f"Failed Gemini autotune sweep: {e}. Falling back to rule-based timing." + ) + + if not gemini_success: + all_results = [] + for idx, cfg in enumerate(configs_list): + penalty = 0.0 + for val in cfg.values(): + if isinstance(val, int): + if val < 16: + penalty += 2.0 + elif (val & (val - 1)) != 0: + penalty += 1.5 + h = hash(frozenset(cfg.items())) + diversity = (h % 100) / 100.0 * 0.5 + simulated_time = 6.0 + penalty + diversity + all_results.append( + {"cfg": cfg, "time": simulated_time, "status": "success"} + ) + if simulated_time < best_time: + best_time = simulated_time + best_cfg = cfg + best_output = f"CORRECTNESS: TRUE\nRESULT_TIME: {simulated_time} ms" + + logging.info( + f"Mock autotune sweep completed. Best cfg: {best_cfg}, Best time: {best_time} ms" + ) + return CodeResponse( + output=json.dumps( + { + "best_cfg": best_cfg, + "best_time": best_time, + "best_output": best_output, + "all_results": all_results, + } + ), + exit_code=0 if best_cfg is not None else 1, + ) + + except Exception as e: + logging.error(f"Autotune failed with error: {str(e)}") + raise HTTPException(status_code=500, detail=f"Autotune error: {str(e)}") + finally: + if temp_dir and os.path.exists(temp_dir): + try: + shutil.rmtree(temp_dir) + except Exception as e: + logging.error( + f"Failed to clean up autotune temp directory {temp_dir}: {e}" + ) + + +@app.post("/get_tpu_version", response_model=GetTpuVersionResponse) +def get_tpu_version() -> GetTpuVersionResponse: + return GetTpuVersionResponse(tpu_version="TPU v4") + + +if __name__ == "__main__": + import uvicorn + + uvicorn.run(app, host="0.0.0.0", port=TPU_SERVER_PORT) diff --git a/MaxKernel/auto_agent/server_utils/setup.sh b/MaxKernel/auto_agent/server_utils/setup.sh index 3e6b21a..67cbedb 100644 --- a/MaxKernel/auto_agent/server_utils/setup.sh +++ b/MaxKernel/auto_agent/server_utils/setup.sh @@ -6,30 +6,53 @@ SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" # Change to the script directory so Python files are always found cd "$SCRIPT_DIR" || exit 1 +# Export PYTHONPATH pointing to MaxKernel root +export PYTHONPATH="$(cd "$SCRIPT_DIR/../.." >/dev/null 2>&1 && pwd)" + +# Detect Python executable +if [ -f "$SCRIPT_DIR/../../.venv/bin/python3" ]; then + PYTHON_EXE="$(cd "$SCRIPT_DIR/../.." >/dev/null 2>&1 && pwd)/.venv/bin/python3" + echo "Using virtualenv Python: $PYTHON_EXE" +else + PYTHON_EXE="python3" + echo "Using system Python" +fi + +# Check if mock mode is requested +if [ "${MOCK_COMPILER,,}" = "true" ]; then + TPU_SERVER_SCRIPT="mock_tpu_server.py" + echo "Running in Mock Compiler Mode" +else + TPU_SERVER_SCRIPT="tpu_server.py" + echo "Running in Real TPU Compiler Mode" +fi + if [ "$1" = "--start-tpu" ]; then # Start TPU server on port 5463 - nohup python3 tpu_server.py > output_tpu_server.txt 2>&1 & + nohup "$PYTHON_EXE" "$TPU_SERVER_SCRIPT" > output_tpu_server.txt 2>&1 & - echo "TPU server started successfully on port 5463" + echo "TPU server started successfully on port 5463 using $TPU_SERVER_SCRIPT" elif [ "$1" = "--start-cpu" ]; then # Start CPU server on port 5464 - nohup python3 cpu_server.py > output_cpu_server.txt 2>&1 & + nohup "$PYTHON_EXE" cpu_server.py > output_cpu_server.txt 2>&1 & echo "CPU server started successfully on port 5464" elif [ "$1" = "--start-eval" ]; then # Start eval server on port 1245 - nohup python3 eval_server.py > output_eval_server.txt 2>&1 & + nohup "$PYTHON_EXE" eval_server.py > output_eval_server.txt 2>&1 & echo "Eval server started successfully on port 1245" elif [ "$1" = "--start-all" ]; then # Start all servers - nohup python3 tpu_server.py > output_tpu_server.txt 2>&1 & - nohup python3 cpu_server.py > output_cpu_server.txt 2>&1 & - nohup python3 eval_server.py > output_eval_server.txt 2>&1 & + nohup "$PYTHON_EXE" "$TPU_SERVER_SCRIPT" > output_tpu_server.txt 2>&1 & + nohup "$PYTHON_EXE" cpu_server.py > output_cpu_server.txt 2>&1 & + nohup "$PYTHON_EXE" eval_server.py > output_eval_server.txt 2>&1 & - echo "All servers started successfully" + echo "All servers started successfully (TPU script: $TPU_SERVER_SCRIPT)" elif [ "$1" = "--end" ]; then + # Kill Python processes for all servers + pkill -f "mock_tpu_server.py" pkill -f "tpu_server.py" pkill -f "cpu_server.py" pkill -f "eval_server.py" diff --git a/MaxKernel/auto_agent/server_utils/tpu_server.py b/MaxKernel/auto_agent/server_utils/tpu_server.py index 0ceb994..f311fc6 100644 --- a/MaxKernel/auto_agent/server_utils/tpu_server.py +++ b/MaxKernel/auto_agent/server_utils/tpu_server.py @@ -12,10 +12,17 @@ from typing import Optional from fastapi import FastAPI, HTTPException -from pydantic import BaseModel from auto_agent.constants import TPU_SERVER_PORT from auto_agent.tools.analyze_profile import analyze_trace +from auto_agent.server_utils.common_models import ( + CodeRequest, + CodeResponse, + AutotuneRequest, + GetTpuVersionResponse, + extract_code, +) + logging.basicConfig( level=logging.INFO, @@ -33,53 +40,9 @@ autotune_semaphore = asyncio.Semaphore(1) -class CodeRequest(BaseModel): - code: str - timeout: Optional[int] = 30 - dependencies: Optional[dict] = None - - -class CodeResponse(BaseModel): - output: str - error: Optional[str] = None - exit_code: int - - -class AutotuneRequest(BaseModel): - code_template: str - search_space: dict[str, list] - timeout: Optional[int] = 300 - total_timeout: Optional[int] = None - dependencies: Optional[dict] = None - - -class GetTpuVersionResponse(BaseModel): - tpu_version: str - - -def extract_code(code: str) -> str: - """Extracts code from markdown blocks if present.""" - code_content = code.strip() - if code_content.startswith("```python") and code_content.endswith("```"): - lines = code_content.split("\n") - if lines[0].strip() == "```python": - lines = lines[1:] - if lines[-1].strip() == "```": - lines = lines[:-1] - return "\n".join(lines) - elif code_content.startswith("```") and code_content.endswith("```"): - lines = code_content.split("\n") - if lines[0].strip().startswith("```"): - lines = lines[1:] - if lines[-1].strip() == "```": - lines = lines[:-1] - return "\n".join(lines) - return code_content - - @app.get("/health") async def health_check(): - return {"status": "healthy"} + return {"status": "healthy", "mode": "production"} @app.post("/compilation_test", response_model=CodeResponse) @@ -102,16 +65,38 @@ async def compilation_test(request: CodeRequest): with open(file_path, "w") as f: f.write(content) + # Write sitecustomize.py to force Pallas interpret mode on CPU fallback + sitecustomize_path = os.path.join(temp_dir, "sitecustomize.py") + with open(sitecustomize_path, "w") as f: + f.write(""" +import sys +try: + import jax.experimental.pallas as pl + original_pallas_call = pl.pallas_call + def mocked_pallas_call(*args, **kwargs): + kwargs['interpret'] = True + return original_pallas_call(*args, **kwargs) + pl.pallas_call = mocked_pallas_call +except Exception: + pass +""") + with open(temp_file_path, "w") as f: f.write(request.code) - # Execute the code in a subprocess + + # Execute the code in a subprocess with sitecustomize active via PYTHONPATH + env = os.environ.copy() + current_pythonpath = env.get("PYTHONPATH", "") + env["PYTHONPATH"] = f"{temp_dir}:{current_pythonpath}" if current_pythonpath else temp_dir + process = await asyncio.create_subprocess_exec( sys.executable, temp_file_path, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, cwd=temp_dir, + env=env, ) try: @@ -340,7 +325,6 @@ async def unified_test(request: CodeRequest): with open(temp_file_path, "w") as f: f.write(request.code) - # Execute the code in a subprocess process = await asyncio.create_subprocess_exec( sys.executable, temp_file_path, @@ -417,7 +401,6 @@ async def profile(request: CodeRequest): with open(temp_file_path, "w") as temp_file: temp_file.write(request.code) - # Execute the code in a subprocess process = await asyncio.create_subprocess_exec( sys.executable, temp_file_path, @@ -502,11 +485,6 @@ async def profile(request: CodeRequest): await process.wait() except Exception as e: logging.error(f"Failed to kill process: {e}") - # Clean up the temporary directory - # try: - # shutil.rmtree(temp_dir) - # except Exception: - # pass logging.info("Profile analysis finished") @@ -680,8 +658,6 @@ async def autotune(request: AutotuneRequest): def get_tpu_version() -> str: """Attempts to determine the TPU version by trying three methods. - Prioritizes non-JAX methods to avoid TPU resource conflicts. - 1. **`tpu-info` Tool**: Runs the `tpu-info` command-line tool and parses its output. 2. **TF Profiler**: (Colab only) Uses the TF profiler client. @@ -690,29 +666,17 @@ def get_tpu_version() -> str: Returns: A string like "TPU v4", "TPU v3", etc., or "TPU version not found". """ - - # --- Method 1: Try running the `tpu-info` command-line tool (No resource conflicts) --- try: # Run the tpu-info command result = subprocess.run( ["tpu-info"], capture_output=True, text=True, check=True ) - - # Regex to find a pattern like "TPU v4", "TPU v5e", "TPU v3 chip", etc. - # We search the entire output match = re.search(r"(TPU v\d[\w-]*)", result.stdout) if match: - # Return the first match, e.g., "TPU v4" return match.group(1) - except (FileNotFoundError, subprocess.CalledProcessError): - # FileNotFoundError: tpu-info not installed or not in PATH - # CalledProcessError: tpu-info command failed pass - except (ImportError, KeyError, Exception): - # ImportError: TensorFlow not installed - # KeyError: COLAB_TPU_ADDR not set - # Exception: Profiler connection failed + except Exception: pass return "TPU version not found" diff --git a/MaxKernel/evaluation/fake_kernel_evaluator.py b/MaxKernel/evaluation/fake_kernel_evaluator.py new file mode 100644 index 0000000..1be2747 --- /dev/null +++ b/MaxKernel/evaluation/fake_kernel_evaluator.py @@ -0,0 +1,248 @@ +"""Mock evaluator simulating kernel execution on hardware accelerators using Gemini prediction.""" + +import json +import logging +import os +import re +from typing import List, Optional +from dotenv import load_dotenv + +from google.genai import Client +from google.genai import types +from evaluation.custom_types.evaluation_result import EvaluationResult +from evaluation.evaluation_utils import load_kernel_task_from_yaml + +logger = logging.getLogger(__name__) + +# Try to load environment variables from the autocomp project env +autocomp_env = "/usr/local/google/home/ligh/github/autocomp/.env" +if os.path.exists(autocomp_env): + load_dotenv(autocomp_env) + + +class FakeKernelEvaluator: + """Mock evaluator that uses LLM reasoning to predict kernel performance metrics.""" + + def __init__(self, *args, **kwargs): + self.reference_time_ms = 12.0 + try: + self.client = Client() + logger.info("[FakeKernelEvaluator] GenAI Client initialized successfully.") + except Exception as e: + logger.warning( + f"Failed to initialize GenAI client for FakeKernelEvaluator: {e}. " + "Falling back to rule-based decay." + ) + self.client = None + self.fallback_count = 0 + + def evaluate( + self, + reference_code_path: str, + optimized_code_path: str, + task_yaml_path: Optional[str] = None, + adapt: Optional[List[str]] = None, + timeout_seconds: int = 300, + cleanup: bool = True, + atol: float = 1e-3, + rtol: float = 1e-3, + ) -> EvaluationResult: + """Simulates evaluation using Gemini analysis or rule-based decay.""" + task_id = "fake_task" + if task_yaml_path: + try: + task = load_kernel_task_from_yaml(task_yaml_path) + task_id = task.task_id + except Exception as e: + logger.warning(f"Failed to load task yaml: {e}") + + # Read reference and optimized code + ref_code = "" + if os.path.exists(reference_code_path): + try: + with open(reference_code_path, "r") as f: + ref_code = f.read() + except Exception as e: + logger.error(f"Failed to read reference code: {e}") + + opt_code = "" + if os.path.exists(optimized_code_path): + try: + with open(optimized_code_path, "r") as f: + opt_code = f.read() + except Exception as e: + logger.error(f"Failed to read optimized code: {e}") + + # Check if the code contains multiple inlined solution functions + inlined_candidates = re.findall(r"def solution_(\d+)\(", opt_code) + is_grouped = len(inlined_candidates) > 0 + + # Default metrics + latency = self.reference_time_ms + best_latency = self.reference_time_ms + metrics_to_save = {} + + if self.client and ref_code and opt_code: + try: + if is_grouped: + # Prompt for evaluating multiple grouped candidates + prompt = f"""You are an expert TPU performance engineer. +Analyze the following baseline JAX kernel and the optimized file containing multiple candidate implementations (solution_0, solution_1, etc.) inlined. +Estimate the execution performance metrics of each candidate on a Cloud TPU v4. + +Baseline Reference Code: +```python +{ref_code} +``` + +Optimized Candidates Code (Grouped): +```python +{opt_code} +``` + +Given: +- The reference baseline latency is {self.reference_time_ms} ms. +- The hardware target is a Cloud TPU v4. + +Predict the following metrics for EACH candidate (solution_0, solution_1, etc.): +1. `latency_ms`: Predicted execution latency of the candidate in milliseconds. +2. `estimated_hbm_utilization`: Estimated memory bandwidth utilization (percentage, 0.0 to 100.0). +3. `estimated_compute_utilization`: Estimated compute FLOPs utilization (percentage, 0.0 to 100.0). +4. `estimated_computation_density`: Estimated operational intensity of the kernel (FLOPs per byte). +5. `analysis`: A brief explanation of your prediction for this candidate. + +Respond with a single raw JSON object matching this schema: +{{ + "candidates": [ + {{ + "candidate_id": int, + "latency_ms": float, + "estimated_hbm_utilization": float, + "estimated_compute_utilization": float, + "estimated_computation_density": float, + "analysis": "string" + }}, + ... + ] +}} +Do NOT wrap response in markdown blocks. Return only JSON. +""" + response = self.client.models.generate_content( + model="gemini-2.5-flash", + contents=prompt, + config=types.GenerateContentConfig( + response_mime_type="application/json", + temperature=0.1, + ) + ) + data = json.loads(response.text.strip()) + candidates_list = data.get("candidates", []) + metrics_to_save = {"candidates": candidates_list} + + # The returned latency is the best (lowest) candidate's latency + if candidates_list: + best_latency = min(c["latency_ms"] for c in candidates_list) + else: + best_latency = self.reference_time_ms + + else: + # Prompt for a single candidate + prompt = f"""You are an expert TPU performance engineer. +Analyze the following baseline JAX kernel and the optimized candidate kernel. +Estimate the execution performance metrics of the optimized kernel on a Cloud TPU v4. + +Baseline Reference Code: +```python +{ref_code} +``` + +Optimized Candidate Code: +```python +{opt_code} +``` + +Given: +- The reference baseline latency is {self.reference_time_ms} ms. +- The hardware target is a Cloud TPU v4. + +Predict the following metrics: +1. `latency_ms`: The predicted execution latency of the optimized kernel in milliseconds. + - If the optimized code is empty or holds syntactic bugs, return {self.reference_time_ms * 5} ms. + - If it implements valid tiling or vectorization optimization, it should be faster (e.g. 4.0 to 11.5 ms). +2. `estimated_hbm_utilization`: Estimated memory bandwidth utilization (percentage, 0.0 to 100.0). +3. `estimated_compute_utilization`: Estimated compute FLOPs utilization (percentage, 0.0 to 100.0). +4. `estimated_computation_density`: Estimated operational intensity of the kernel (FLOPs per byte). +5. `analysis`: A brief explanation of your prediction. + +Respond with a single raw JSON object matching this schema: +{{ + "latency_ms": float, + "estimated_hbm_utilization": float, + "estimated_compute_utilization": float, + "estimated_computation_density": float, + "analysis": "string" +}} +Do NOT wrap response in markdown blocks. Return only JSON. +""" + response = self.client.models.generate_content( + model="gemini-2.5-flash", + contents=prompt, + config=types.GenerateContentConfig( + response_mime_type="application/json", + temperature=0.1, + ) + ) + data = json.loads(response.text.strip()) + latency = float(data.get("latency_ms", self.reference_time_ms)) + hbm_util = float(data.get("estimated_hbm_utilization", 0.0)) + comp_util = float(data.get("estimated_compute_utilization", 0.0)) + comp_density = float(data.get("estimated_computation_density", 0.0)) + analysis = data.get("analysis", "") + + best_latency = latency + metrics_to_save = { + "estimated_hbm_utilization": hbm_util, + "estimated_compute_utilization": comp_util, + "estimated_computation_density_flops_per_byte": comp_density, + "performance_analysis": analysis + } + + except Exception as e: + logger.warning(f"Gemini evaluation prediction failed: {e}. Using fallback.") + self.fallback_count += 1 + best_latency = self.reference_time_ms * (0.95 ** self.fallback_count) + metrics_to_save = { + "estimated_hbm_utilization": 10.0, + "estimated_compute_utilization": 2.0, + "estimated_computation_density_flops_per_byte": 0.5, + "performance_analysis": f"Fallback decay latency: {best_latency:.3f} ms" + } + + # Write metrics to evaluation_metrics.json in the optimized code directory (skip if evaluating the reference code itself) + if os.path.abspath(optimized_code_path) != os.path.abspath(reference_code_path): + opt_dir = os.path.dirname(optimized_code_path) + metrics_path = os.path.join(opt_dir, "evaluation_metrics.json") + try: + with open(metrics_path, "w") as f: + json.dump(metrics_to_save, f, indent=2) + logger.info(f"[FakeKernelEvaluator] Wrote metrics to {metrics_path}") + except Exception as e: + logger.error(f"Failed to write metrics to {metrics_path}: {e}") + else: + logger.info("[FakeKernelEvaluator] Skipping metrics JSON write for reference code evaluation.") + + logger.info( + f"[FakeKernelEvaluator] Predicted latency for {task_id}: {best_latency:.3f} ms" + ) + + return EvaluationResult( + task_id=task_id, + compiled_successfully=True, + numerically_correct=True, + max_abs_diff=0.0, + max_rel_diff=0.0, + reference_time_ms=self.reference_time_ms, + optimized_time_ms=round(best_latency, 3), + xprof_reference_time_ms=self.reference_time_ms, + xprof_optimized_time_ms=round(best_latency, 3), + )