Prepares (modifies in-place) the graph to be served to the front-end. For now, it supports filtering out attributes that are too large to be shown in the graph UI. Args: graph: The GraphDef proto message. limit_attr_size: Maximum allowed size in bytes, before the attribute
(
graph, limit_attr_size=1024, large_attrs_key="_too_large_attrs"
)
| 19 | |
| 20 | |
| 21 | def prepare_graph_for_ui( |
| 22 | graph, limit_attr_size=1024, large_attrs_key="_too_large_attrs" |
| 23 | ): |
| 24 | """Prepares (modifies in-place) the graph to be served to the front-end. |
| 25 | |
| 26 | For now, it supports filtering out attributes that are |
| 27 | too large to be shown in the graph UI. |
| 28 | |
| 29 | Args: |
| 30 | graph: The GraphDef proto message. |
| 31 | limit_attr_size: Maximum allowed size in bytes, before the attribute |
| 32 | is considered large. Default is 1024 (1KB). Must be > 0 or None. |
| 33 | If None, there will be no filtering. |
| 34 | large_attrs_key: The attribute key that will be used for storing attributes |
| 35 | that are too large. Default is '_too_large_attrs'. Must be != None if |
| 36 | `limit_attr_size` is != None. |
| 37 | |
| 38 | Raises: |
| 39 | ValueError: If `large_attrs_key is None` while `limit_attr_size != None`. |
| 40 | ValueError: If `limit_attr_size` is defined, but <= 0. |
| 41 | """ |
| 42 | # TODO(@davidsoergel): detect whether a graph has been filtered already |
| 43 | # (to a limit_attr_size <= what is requested here). If it is already |
| 44 | # filtered, return immediately. |
| 45 | |
| 46 | # Check input for validity. |
| 47 | if limit_attr_size is not None: |
| 48 | if large_attrs_key is None: |
| 49 | raise ValueError( |
| 50 | "large_attrs_key must be != None when limit_attr_size" |
| 51 | "!= None." |
| 52 | ) |
| 53 | |
| 54 | if limit_attr_size <= 0: |
| 55 | raise ValueError( |
| 56 | "limit_attr_size must be > 0, but is %d" % limit_attr_size |
| 57 | ) |
| 58 | |
| 59 | # Filter only if a limit size is defined. |
| 60 | if limit_attr_size is not None: |
| 61 | for node in graph.node: |
| 62 | # Go through all the attributes and filter out ones bigger than the |
| 63 | # limit. |
| 64 | keys = list(node.attr.keys()) |
| 65 | for key in keys: |
| 66 | size = node.attr[key].ByteSize() |
| 67 | if size > limit_attr_size or size < 0: |
| 68 | del node.attr[key] |
| 69 | # Add the attribute key to the list of "too large" attributes. |
| 70 | # This is used in the info card in the graph UI to show the user |
| 71 | # that some attributes are too large to be shown. |
| 72 | node.attr[large_attrs_key].list.s.append( |
| 73 | key.encode("utf-8") |
| 74 | ) |