From 7099af445075fc9fd06c373d3cc68eaa7d5f0d98 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Mon, 30 Sep 2024 09:35:36 +0800 Subject: [PATCH] VIZ show rendering errors (#6807) * VIZ show rendering errors * show the entire traceback --- viz/serve.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/viz/serve.py b/viz/serve.py index bfd1b7dbee..dd14310a10 100755 --- a/viz/serve.py +++ b/viz/serve.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 from __future__ import annotations from typing import Dict, List, Optional, Tuple -import pickle, os, sys, time, threading, webbrowser, json, difflib, contextlib, re +import pickle, os, sys, time, threading, webbrowser, json, difflib, contextlib, re, traceback from dataclasses import dataclass, asdict from urllib.parse import parse_qs, urlparse from http.server import HTTPServer, BaseHTTPRequestHandler @@ -96,10 +96,12 @@ def load_kernels(contexts:List[TrackedRewriteContext]) -> List[KernelRet]: kernel_name = "" code = "" for ctx in contexts: - if ctx.loc[0].split("/")[-1] == "schedule.py": - si_ctx = ScheduleItemContext(bufs=tuple(x.arg for x in ctx.sink.sparents if x.op is UOps.BUFFER)) - with Context(TRACK_MATCH_STATS=0): kernel_name, code = (prg:=get_runner(Device.DEFAULT, full_ast_rewrite(ctx.sink, si_ctx)).p).name, prg.src - elif ctx.kernel_name is not None: kernel_name, code = ctx.kernel_name, "" + try: + if ctx.loc[0].split("/")[-1] == "schedule.py": + si_ctx = ScheduleItemContext(bufs=tuple(x.arg for x in ctx.sink.sparents if x.op is UOps.BUFFER)) + with Context(TRACK_MATCH_STATS=0): kernel_name, code = (prg:=get_runner(Device.DEFAULT, full_ast_rewrite(ctx.sink, si_ctx)).p).name, prg.src + elif ctx.kernel_name is not None: kernel_name, code = ctx.kernel_name, "" + except Exception: kernel_name, code = "RENDERING_ERROR", traceback.format_exc() if ret.get(k:=to_function_name(kernel_name)) is None: ret[k] = KernelRet(k, code, []) ret[k].ctxs.append(ctx) return list(ret.values())