From 568a4b5d92b642f98dd6e17b3bb78407ffbb8f03 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Wed, 16 Oct 2024 15:14:47 +0300
Subject: [PATCH] viz touchups [pr] (#7095)

* viz touchups [pr]

* check if port busy

* url
---
 tinygrad/viz/eslint.config.mjs |  9 ---------
 tinygrad/viz/index.html        | 31 +++++++++++++------------------
 tinygrad/viz/serve.py          | 11 +++++++----
 3 files changed, 20 insertions(+), 31 deletions(-)
 delete mode 100644 tinygrad/viz/eslint.config.mjs

diff --git a/tinygrad/viz/eslint.config.mjs b/tinygrad/viz/eslint.config.mjs
deleted file mode 100644
index d95bcaf69177..000000000000
--- a/tinygrad/viz/eslint.config.mjs
+++ /dev/null
@@ -1,9 +0,0 @@
-import globals from "globals";
-import pluginJs from "@eslint/js";
-import pluginHtml from "eslint-plugin-html";
-
-export default [
-  {files: ["**/*.html"], plugins: {html: pluginHtml}, rules: {"max-len": ["error", {"code": 150}]}},
-  {languageOptions: {globals: globals.browser}},
-  pluginJs.configs.recommended,
-];
diff --git a/tinygrad/viz/index.html b/tinygrad/viz/index.html
index 52c608815522..c6de5ebfefbf 100644
--- a/tinygrad/viz/index.html
+++ b/tinygrad/viz/index.html
@@ -58,7 +58,7 @@
   }
   .node rect {
     stroke: #4a4b57;
-    stroke-width: 1.5px;
+    stroke-width: 1.4px;
     rx: 8px;
     ry: 8px;
   }
@@ -81,6 +81,9 @@
     width: 100%;
     height: 100%;
   }
+  .main-container > * + * {
+    margin-left: 8px;
+  }
   .container {
     background-color: #0f1018;
     height: 100%;
@@ -88,22 +91,13 @@
     padding: 8px;
     position: relative;
   }
-  .container > * + * {
+  .container > * + *, .rewrite-container > * + * {
     margin-top: 12px;
   }
-  .rewrite-container > * + * {
-    margin-top: 12px;
-  }
-  .main-container > * + * {
-    margin-left: 8px;
-  }
   .kernel-list {
     width: 10%;
     overflow-y: auto;
   }
-  .kernel-list > ul > ul {
-    padding-left: 4px;
-  }
   .kernel-list > ul > * + * {
     margin-top: 4px;
   }
@@ -153,7 +147,8 @@
   </div>
   <script>
   /* global hljs, dagreD3, d3, DOMPurify */
-  // extra definitions for UOps
+
+  // **** hljs extra definitions for UOps and float4
   hljs.registerLanguage("python", (hljs) => ({
     ...hljs.getLanguage("python"),
     case_insensitive: false,
@@ -165,12 +160,12 @@
       ...hljs.getLanguage("python").contains,
     ]
   }));
-  // extra definitions for float4
   hljs.registerLanguage("cpp", (hljs) => ({
     ...hljs.getLanguage('cpp'),
     contains: [{ begin: '\\b(?:float|half)[0-9]+\\b', className: 'type' }, ...hljs.getLanguage('cpp').contains]
   }));
 
+  // **** D3
   function recenterRects(svg, zoom) {
     const svgBounds = svg.node().getBoundingClientRect();
     for (const rect of svg.node().querySelectorAll("rect")) {
@@ -182,7 +177,6 @@
     }
     svg.call(zoom.transform, d3.zoomIdentity)
   }
-
   function renderGraph(graph, additions) {
     const g = new dagreD3.graphlib.Graph({ compound: true }).setGraph({ rankdir: "LR" }).setDefaultEdgeLabel(function() { return {}; });
     g.setNode("addition", {label: "", clusterLabelPos: "top", style: additions.length !== 0 ? "fill: rgba(26, 27, 38, 0.5);" : "display: none;"});
@@ -209,11 +203,10 @@
     render(inner, g);
   }
 
-  function toPath(loc) {
-    const [fp, lineno] = loc;
-    return `${fp.split("/").pop()}:${lineno}`
-  }
+  // **** extra helpers
+  const toPath = ([fp, lineno]) => `${fp.split("/").pop()}:${lineno}`;
 
+  // **** main loop
   var ret = {};
   var cache = {};
   var kernels = null;
@@ -359,6 +352,8 @@
       metadata.appendChild(Object.assign(document.createElement("p"), { textContent: `No rewrites in ${toPath(ret.loc)}.` }));
     }
   }
+
+  // **** keyboard shortcuts
   document.addEventListener("keydown", async function(event) {
     // up and down change the UOp or kernel from the list
     if (!expandKernel) {
diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py
index 152e57021840..9ac1684828ee 100755
--- a/tinygrad/viz/serve.py
+++ b/tinygrad/viz/serve.py
@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-import multiprocessing, pickle, functools, difflib, os, threading, json, time, sys, webbrowser
+import multiprocessing, pickle, functools, difflib, os, threading, json, time, sys, webbrowser, socket
 from http.server import HTTPServer, BaseHTTPRequestHandler
 from urllib.parse import parse_qs, urlparse
 from dataclasses import asdict, dataclass
@@ -122,6 +122,9 @@ def reloader():
     time.sleep(0.1)
 
 if __name__ == "__main__":
+  with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+    if s.connect_ex(((HOST:="http://127.0.0.1").replace("http://", ""), PORT:=getenv("PORT", 8000))) == 0:
+      raise RuntimeError(f"{HOST}:{PORT} is occupied! use PORT= to change.")
   stop_reloader = threading.Event()
   multiprocessing.current_process().name = "VizProcess"    # disallow opening of devices
   st = time.perf_counter()
@@ -133,12 +136,12 @@ def reloader():
     ret = [get_details(*args) for v in tqdm(kernels) for args in v]
     print(f"fuzzed {len(ret)} rewrite details")
   print("*** loaded kernels")
-  server = HTTPServer(('', PORT:=getenv("PORT", 8000)), Handler)
+  server = HTTPServer(('', PORT), Handler)
   reloader_thread = threading.Thread(target=reloader)
   reloader_thread.start()
-  print(f"*** started viz on http://127.0.0.1:{PORT}")
+  print(f"*** started viz on {HOST}:{PORT}")
   print(colored(f"*** ready in {(time.perf_counter()-st)*1e3:4.2f}ms", "green"))
-  if getenv("BROWSER", 0): webbrowser.open(f"http://127.0.0.1:{PORT}")
+  if getenv("BROWSER", 0): webbrowser.open(f"{HOST}:{PORT}")
   try: server.serve_forever()
   except KeyboardInterrupt:
     print("*** viz is shutting down...")