lablab-ai-amd-developer-hackathon/Sentinel-ROCm
2
1import gradio as gr 2import os3from typing import TypedDict4from langgraph.graph import StateGraph, START, END5from langchain_openai import ChatOpenAI6 7class MigrationState(TypedDict):8 codigo_original: str9 codigo_traducido: str10 analisis_seguridad: str11 12# Nota: El código trabaja actualmente con el modelo DeepSeek V3.1. 13# Sin embargo, se recomienda trabajar con Qwen2.5-Coder 32B para un mejor análisis de código. 14llm = ChatOpenAI(15 api_key = os.environ.get("API_KEY"),16 base_url="https://api.fireworks.ai/inference/v1",17 model="accounts/fireworks/models/deepseek-v3p1" 18)19 20def nodo_traductor(state: MigrationState):21 prompt = f"Actúa como un experto en AMD ROCm. Traduce este código NVIDIA CUDA a AMD HIP. Solo devuelve el código final, sin explicaciones:\n\n{state['codigo_original']}"22 respuesta = llm.invoke(prompt)23 24 codigo_limpio = respuesta.content25 if codigo_limpio.startswith("```"):26 lineas = codigo_limpio.split("\n")27 codigo_limpio = "\n".join(lineas[1:-1])28 29 return {"codigo_traducido": codigo_limpio}30 31def nodo_auditor(state: MigrationState):32 prompt = f"Revisa este código AMD HIP. Confirma si la migración fue correcta y advierte sobre posibles fugas de memoria. Sé breve (máximo 3 líneas):\n\n{state['codigo_traducido']}"33 respuesta = llm.invoke(prompt)34 return {"analisis_seguridad": respuesta.content}35 36workflow = StateGraph(MigrationState)37workflow.add_node("traductor", nodo_traductor)38workflow.add_node("auditor", nodo_auditor)39workflow.add_edge(START, "traductor")40workflow.add_edge("traductor", "auditor")41workflow.add_edge("auditor", END)42 43agente_rocm = workflow.compile()44 45DUMMY_CUDA_CODE = """#include <iostream>46 47__global__ void vectorAdd(const float *A, const float *B, float *C, int numElements) {48 int i = blockDim.x * blockIdx.x + threadIdx.x;49 if (i < numElements) {50 C[i] = A[i] + B[i];51 }52}53 54int main() {55 return 0;56}57"""58 59def leer_archivo(archivo):60 if archivo is not None:61 try:62 with open(archivo.name, "r", encoding="utf-8") as f:63 return f.read()64 except Exception as e:65 return f"Error al leer el archivo: {str(e)}"66 return ""67 68def chat_con_agente(mensaje_usuario, historial):69 if not historial:70 historial = []71 72 prompt_chat = f"""Eres el Ingeniero Líder de ROCm-Sentinel. 73 El usuario te da esta instrucción para la migración: '{mensaje_usuario}'. 74 Responde en una frase corta y técnica cómo integrarás esta instrucción en el proceso de traducción a AMD HIP."""75 76 try:77 respuesta = llm.invoke(prompt_chat)78 historial.append({"role": "user", "content": mensaje_usuario})79 historial.append({"role": "assistant", "content": respuesta.content})80 except Exception:81 historial.append({"role": "user", "content": mensaje_usuario})82 historial.append({"role": "assistant", "content": "Entendido. Aplicaré esas optimizaciones técnicas al código."})83 84 return "", historial85 86def simular_migracion(cuda_code, historial):87 if not historial: historial = []88 if not cuda_code:89 yield "Error: Ingresa código.", historial, "Desconocido", ""90 return91 92 historial.append({"role": "assistant", "content": "🚀 **Proceso iniciado.** El Agente Traductor está analizando los kernels de CUDA..."})93 yield "Estado: Analizando Kernels...", historial, "Analizando...", ""94 95 try:96 resultado = agente_rocm.invoke({"codigo_original": cuda_code})97 rocm_code = resultado["codigo_traducido"]98 reporte_auditor = resultado["analisis_seguridad"]99 100 historial.append({"role": "assistant", "content": f"✅ **Migración Exitosa.**\n\n🛡️ **Auditoría de Seguridad:**\n{reporte_auditor}"})101 yield "Estado: Completado con éxito", historial, "NVIDIA CUDA (nvcc) -> AMD HIP (hipcc)", rocm_code102 103 except Exception as e:104 historial.append({"role": "assistant", "content": f"❌ Error en el flujo: {str(e)}"})105 yield "Estado: Error crítico", historial, "Error", ""106 107with gr.Blocks() as demo:108 gr.Markdown(109 """110 <div style='text-align: center;'>111 <h1 style='color: #FF3232;'>ROCm-Sentinel</h1>112 <p><strong>Advanced CUDA to AMD ROCm Migration Agent</strong></p>113 </div>114 """115 )116 117 with gr.Row():118 with gr.Column(scale=1):119 archivo_subida = gr.File(120 label="📁 Subir archivo fuente", 121 file_types=[".cu", ".cpp", ".txt", ".h"]122 )123 gr.Markdown("### 📜 Código Origen (CUDA)")124 input_code = gr.Code(125 label="Editor de entrada", 126 language="cpp", 127 lines=15, 128 value=DUMMY_CUDA_CODE129 )130 131 with gr.Column(scale=1):132 lenguaje_detectado = gr.Textbox(133 label="🔍 Firma / Lenguaje Detectado", 134 interactive=False135 )136 estado_migracion = gr.Textbox(137 label="⚡ Monitor de Proceso (Estado)",138 interactive=False,139 placeholder="Listo para procesar..."140 )141 gr.Markdown("### 🛠️ Código Migrado (AMD HIP)")142 output_code = gr.Code(143 label="Resultado de la migración", 144 language="cpp", 145 lines=15, 146 interactive=False147 )148 149 with gr.Row():150 migrate_btn = gr.Button(151 "🚀 INICIAR MIGRACIÓN SEGURA A ROCm", 152 variant="primary", 153 size="lg"154 )155 156 with gr.Row():157 with gr.Column(scale=1):158 gr.Markdown("### 🤖 Centro de Control del Agente")159 chatbot = gr.Chatbot(160 label="Interacción con el Agente de IA", 161 height=300162 )163 chat_input = gr.Textbox(164 label="Instrucciones adicionales para el Agente", 165 placeholder="Ej: 'Usa memoria unificada' o 'Optimiza para hilos específicos'..."166 )167 168 archivo_subida.change(fn=leer_archivo, inputs=archivo_subida, outputs=input_code)169 170 chat_input.submit(171 fn=chat_con_agente, 172 inputs=[chat_input, chatbot], 173 outputs=[chat_input, chatbot]174 )175 176 migrate_btn.click(177 fn=simular_migracion, 178 inputs=[input_code, chatbot], 179 outputs=[estado_migracion, chatbot, lenguaje_detectado, output_code]180 )181 182if __name__ == "__main__":183 demo.launch(theme=gr.themes.Monochrome())