Fix bugs in NAS UI (deferred from previous releases) (#2552)

* Fix NASUI bugs

* Add pipeline for NASUI

* update log

* Move graph utils test to test folder

src/nasui/package.json

@@ -20,6 +20,7 @@
     "start": "react-scripts start",
     "build": "react-scripts build",
     "eject": "react-scripts eject",
+    "test": "react-scripts test",
     "backend": "node server.js"
   },
   "eslintConfig": {

src/nasui/src/App.tsx

@@ -29,6 +29,7 @@ import ExpandMoreIcon from '@material-ui/icons/ExpandMore';
 import List from '@material-ui/core/List';
 import ListItem from '@material-ui/core/ListItem';
 import Backdrop from '@material-ui/core/Backdrop';
+import Tooltip from '@material-ui/core/Tooltip';
 import Chart from './Chart';
 import { Graph } from './graphUtils';
 
@@ -272,15 +273,21 @@ class App extends React.Component<AppProps, AppState>  {
             <Typography variant='h6' className={classes.title}>
               NNI NAS Board
             </Typography>
-            <IconButton color='inherit' onClick={handleLayoutStateChanged(true)}>
-              <ShuffleIcon />
-            </IconButton>
-            <IconButton color='inherit' onClick={this.refresh}>
-              <RefreshIcon />
-            </IconButton>
-            <IconButton color='inherit' onClick={handleSettingsDialogToggle(true)}>
-              <SettingsIcon />
-            </IconButton>
+            <Tooltip title="Re-layout graph">
+              <IconButton color='inherit' onClick={handleLayoutStateChanged(true)}>
+                <ShuffleIcon />
+              </IconButton>
+            </Tooltip>
+            <Tooltip title="Refresh">
+              <IconButton color='inherit' onClick={this.refresh}>
+                <RefreshIcon />
+              </IconButton>
+            </Tooltip>
+            <Tooltip title="Settings">
+              <IconButton color='inherit' onClick={handleSettingsDialogToggle(true)}>
+                <SettingsIcon />
+              </IconButton>
+            </Tooltip>
           </Toolbar>
         </AppBar>
         <AppBar position='fixed' color='default' className={classes.bottomAppBar}>

src/nasui/src/Chart.tsx

@@ -91,6 +91,7 @@ export default class Chart extends React.Component<ChartProps> {
   expandSet: Set<string> = new Set<string>();
   elemWeight: Map<string, number> = new Map<string, number>();
   graphEl: any[] = [];
+  private firstUpdate = true;
 
   componentDidMount() {
     this.cyInstance = cytoscape({
@@ -149,7 +150,8 @@ export default class Chart extends React.Component<ChartProps> {
       }
     } else {
       // re-calculate collapse
-      this.renderGraph(true);
+      this.renderGraph(true, this.firstUpdate);
+      this.firstUpdate = false;
     }
   }
 
@@ -245,7 +247,7 @@ export default class Chart extends React.Component<ChartProps> {
     setTimeout(_render, 100);
   }
 
-  private renderGraph(graphChanged: boolean) {
+  private renderGraph(graphChanged: boolean, fit: boolean = false) {
     const { graph } = this.props;
     if (graph === undefined)
       return;
@@ -269,7 +271,9 @@ export default class Chart extends React.Component<ChartProps> {
           this.cyInstance!.json({
             elements: graphEl
           });
-          layout.fit = false;
+          if (!fit) {
+            layout.fit = false;
+          }
           eles.layout(layout).run();
         }
       } else {

src/nasui/src/__tests__/graphUtils.ts

+import { Graph } from '../graphUtils';
+import *  as fs from 'fs';
+
+describe('Graph utils test', () => {
+  it('Graph constructor darts', () => {
+    const graph = new Graph(JSON.parse(fs.readFileSync('assets/darts/graph.json').toString()), true);
+    const activation = JSON.parse(fs.readFileSync('assets/darts/log').toString().split('\n')[0]);
+    expect(graph.nodes.length).toEqual(1842);
+    expect(graph.edges.length).toEqual(927);
+    const weights = graph.weightFromMutables(activation);
+    expect(weights.get('["CNN/ModuleList[cells]/Cell[1]/ModuleList[mutable_ops]/Node[0]/InputChoice[input_switch]/input.228",' +
+                       '"CNN/ModuleList[cells]/Cell[1]/ModuleList[mutable_ops]/Node[1]/ModuleList[ops]/' + 
+                       'LayerChoice[2]/PoolBN[maxpool]/MaxPool2d[pool]/input.229"]')).toBeCloseTo(0.125, 3);
+  });
+
+  it('Graph constructor naive', () => {
+    const graph = new Graph(JSON.parse(fs.readFileSync('assets/naive/graph.json').toString()), true);
+    expect(graph.nodes.length).toEqual(51);
+    expect(graph.edges.length).toEqual(37);
+    expect(graph.mutableEdges.get('LayerChoice1')![0].length).toEqual(5);
+    expect(graph.mutableEdges.get('LayerChoice1')![1].length).toEqual(5);
+    expect(graph.mutableEdges.get('LayerChoice2')![0].length).toEqual(5);
+    expect(graph.mutableEdges.get('LayerChoice2')![1].length).toEqual(5);
+    expect(graph.mutableEdges.get('InputChoice3')![0].length).toEqual(4);
+  });
+});

src/nasui/src/graphUtils.ts

@@ -47,6 +47,10 @@ export class NodeTs {
   isParent(): boolean {
     return this.children.length > 0;
   }
+
+  toString(): string {
+    return `Node(id=${this.id})`;
+  }
 };
 
 export class Edge {
@@ -59,6 +63,10 @@ export class Edge {
     this.target = target;
     this.id = JSON.stringify([this.source.id, this.target.id]);
   }
+
+  toString(): string {
+    return `Edge(${this.source} -> ${this.target})`;
+  }
 };
 
 interface NodeSummary {
@@ -76,12 +84,12 @@ export class Graph {
   nodes: NodeTs[];
   edges: Edge[];
   defaultExpandSet: Set<string>;
+  mutableEdges: Map<string, Edge[][]>;
   private id2idx: Map<string, number>;
   private edgeId2idx: Map<string, number>;
   private forwardGraph: Map<string, string[]>;
   private backwardGraph: Map<string, string[]>;
   private node2edge: Map<string, Edge[]>;
-  private mutableEdges: Map<string, Edge[][]>;
 
   private build() {
     this.id2idx.clear();
@@ -359,17 +367,25 @@ export class Graph {
   weightFromMutables(mutable: any): Map<string, number> {
     const elemWeight = new Map<string, number>();
     Object.entries(mutable).forEach(entry => {
+      const elemWeightPartial = new Map<string, number>();
       const key = entry[0];
       const weights = entry[1] as number[];
       this.mutableEdges.get(key)!.forEach((edges: any, i: number) => {
         edges.forEach((edge: any) => {
-          if (elemWeight.has(edge.id)) {
-            elemWeight.set(edge.id, elemWeight.get(edge.id)! + weights[i]);
+          if (elemWeightPartial.has(edge.id)) {
+            elemWeightPartial.set(edge.id, elemWeightPartial.get(edge.id)! + weights[i]);
           } else {
-            elemWeight.set(edge.id, weights[i]);
+            elemWeightPartial.set(edge.id, weights[i]);
           }
         })
       });
+      elemWeightPartial.forEach((v, k) => {
+        if (elemWeight.has(k)) {
+          elemWeight.set(k, Math.min(elemWeight.get(k)!, v));
+        } else {
+          elemWeight.set(k, v);
+        }
+      });
     });
     this.nodes.forEach(node => {
       const edges = this.connectedEdges(node.id);

test/scripts/unittest.sh

@@ -26,6 +26,12 @@ echo ""
 echo "===========================Testing: nni_manager==========================="
 npm run test
 
+# -------------For NASUI unittest-------------
+cd ${CWD}/../src/nasui
+echo ""
+echo "===========================Testing: nasui==========================="
+CI=true npm test
+
 ## ------Run nnictl unit test------
 echo ""
 echo "===========================Testing: nnictl==========================="