Run flake8 in Travis and make code PEP8 compliant. (#387)

webui/backend/ray_ui.py

@@ -6,17 +6,17 @@ import collections
 import datetime
 import json
 import numpy as np
-import os
-import redis
-import sys
 import time
 import websockets
 
 # Import flatbuffer bindings.
-from ray.core.generated.LocalSchedulerInfoMessage import LocalSchedulerInfoMessage
+from ray.core.generated.LocalSchedulerInfoMessage import \
+    LocalSchedulerInfoMessage
 
-parser = argparse.ArgumentParser(description="parse information for the web ui")
-parser.add_argument("--redis-address", required=True, type=str, help="the address to use for redis")
+parser = argparse.ArgumentParser(
+    description="parse information for the web ui")
+parser.add_argument("--redis-address", required=True, type=str,
+                    help="the address to use for redis")
 
 loop = asyncio.get_event_loop()
 
@@ -25,27 +25,36 @@ IDENTIFIER_LENGTH = 20
 # This prefix must match the value defined in ray_redis_module.cc.
 DB_CLIENT_PREFIX = b"CL:"
 
+
 def hex_identifier(identifier):
   return binascii.hexlify(identifier).decode()
 
+
 def identifier(hex_identifier):
   return binascii.unhexlify(hex_identifier)
 
+
 def key_to_hex_identifier(key):
-  return hex_identifier(key[(key.index(b":") + 1):(key.index(b":") + IDENTIFIER_LENGTH + 1)])
+  return hex_identifier(
+      key[(key.index(b":") + 1):(key.index(b":") + IDENTIFIER_LENGTH + 1)])
+
 
 def timestamp_to_date_string(timestamp):
   """Convert a time stamp returned by time.time() to a formatted string."""
-  return datetime.datetime.fromtimestamp(timestamp).strftime("%Y/%m/%d %H:%M:%S")
+  return (datetime.datetime.fromtimestamp(timestamp)
+          .strftime("%Y/%m/%d %H:%M:%S"))
+
 
 def key_to_hex_identifiers(key):
-  # Extract worker_id and task_id from key of the form prefix:worker_id:task_id.
+  # Extract worker_id and task_id from key of the form
+  # prefix:worker_id:task_id.
   offset = key.index(b":") + 1
   worker_id = hex_identifier(key[offset:(offset + IDENTIFIER_LENGTH)])
   offset += IDENTIFIER_LENGTH + 1
   task_id = hex_identifier(key[offset:(offset + IDENTIFIER_LENGTH)])
   return worker_id, task_id
 
+
 async def hgetall_as_dict(redis_conn, key):
   fields = await redis_conn.execute("hgetall", key)
   return {fields[2 * i]: fields[2 * i + 1] for i in range(len(fields) // 2)}
@@ -55,6 +64,7 @@ async def hgetall_as_dict(redis_conn, key):
 local_schedulers = {}
 errors = []
 
+
 def duration_to_string(duration):
   """Format a duration in seconds as a string.
 
@@ -79,8 +89,10 @@ def duration_to_string(duration):
     duration_str = "{} microseconds".format(int(duration * 1000000))
   return duration_str
 
+
 async def handle_get_statistics(websocket, redis_conn):
-  cluster_start_time = float(await redis_conn.execute("get", "redis_start_time"))
+  cluster_start_time = float(await redis_conn.execute("get",
+                                                      "redis_start_time"))
   start_date = timestamp_to_date_string(cluster_start_time)
 
   uptime = duration_to_string(time.time() - cluster_start_time)
@@ -90,7 +102,9 @@ async def handle_get_statistics(websocket, redis_conn):
   for client_key in client_keys:
     client_fields = await hgetall_as_dict(redis_conn, client_key)
     clients.append(client_fields)
-  ip_addresses = list(set([client[b"node_ip_address"].decode("ascii") for client in clients if client[b"client_type"] == b"local_scheduler"]))
+  ip_addresses = list(set([client[b"node_ip_address"].decode("ascii")
+                           for client in clients
+                           if client[b"client_type"] == b"local_scheduler"]))
   num_nodes = len(ip_addresses)
   reply = {"uptime": uptime,
            "start_date": start_date,
@@ -98,18 +112,22 @@ async def handle_get_statistics(websocket, redis_conn):
            "addresses": ip_addresses}
   await websocket.send(json.dumps(reply))
 
+
 async def handle_get_drivers(websocket, redis_conn):
   keys = await redis_conn.execute("keys", "Drivers:*")
   drivers = []
   for key in keys:
     driver_fields = await hgetall_as_dict(redis_conn, key)
-    driver_info = {"node ip address": driver_fields[b"node_ip_address"].decode("ascii"),
-                   "name": driver_fields[b"name"].decode("ascii")}
+    driver_info = {
+        "node ip address": driver_fields[b"node_ip_address"].decode("ascii"),
+        "name": driver_fields[b"name"].decode("ascii")}
 
-    driver_info["start time"] = timestamp_to_date_string(float(driver_fields[b"start_time"]))
+    driver_info["start time"] = timestamp_to_date_string(
+        float(driver_fields[b"start_time"]))
 
     if b"end_time" in driver_fields:
-      duration = float(driver_fields[b"end_time"]) - float(driver_fields[b"start_time"])
+      duration = (float(driver_fields[b"end_time"]) -
+                  float(driver_fields[b"start_time"]))
     else:
       duration = time.time() - float(driver_fields[b"start_time"])
     driver_info["duration"] = duration_to_string(duration)
@@ -129,17 +147,20 @@ async def handle_get_drivers(websocket, redis_conn):
   reply = sorted(drivers, key=(lambda driver: driver["start time"]))[::-1]
   await websocket.send(json.dumps(reply))
 
+
 async def listen_for_errors(redis_ip_address, redis_port):
-  pubsub_conn = await aioredis.create_connection((redis_ip_address, redis_port), loop=loop)
-  data_conn = await aioredis.create_connection((redis_ip_address, redis_port), loop=loop)
+  pubsub_conn = await aioredis.create_connection(
+      (redis_ip_address, redis_port), loop=loop)
+  data_conn = await aioredis.create_connection((redis_ip_address, redis_port),
+                                               loop=loop)
 
   error_pattern = "__keyspace@0__:ErrorKeys"
-  psub = await pubsub_conn.execute_pubsub("psubscribe", error_pattern)
+  await pubsub_conn.execute_pubsub("psubscribe", error_pattern)
   channel = pubsub_conn.pubsub_patterns[error_pattern]
   print("Listening for error messages...")
   index = 0
   while (await channel.wait_message()):
-    msg = await channel.get()
+    await channel.get()
     info = await data_conn.execute("lrange", "ErrorKeys", index, -1)
 
     for error_key in info:
@@ -154,6 +175,7 @@ async def listen_for_errors(redis_ip_address, redis_port):
                      "error": result})
       index += 1
 
+
 async def handle_get_errors(websocket):
   """Send error messages to the frontend."""
   await websocket.send(json.dumps(errors))
@@ -161,6 +183,7 @@ async def handle_get_errors(websocket):
 node_info = collections.OrderedDict()
 worker_info = collections.OrderedDict()
 
+
 async def handle_get_recent_tasks(websocket, redis_conn, num_tasks):
   # First update the cache of worker information.
   worker_keys = await redis_conn.execute("keys", "Workers:*")
@@ -168,7 +191,8 @@ async def handle_get_recent_tasks(websocket, redis_conn, num_tasks):
     worker_id = hex_identifier(key[len("Workers:"):])
     if worker_id not in worker_info:
       worker_info[worker_id] = await hgetall_as_dict(redis_conn, key)
-      node_ip_address = worker_info[worker_id][b"node_ip_address"].decode("ascii")
+      node_ip_address = (worker_info[worker_id][b"node_ip_address"]
+                         .decode("ascii"))
       if node_ip_address not in node_info:
         node_info[node_ip_address] = {"workers": []}
       node_info[node_ip_address]["workers"].append(worker_id)
@@ -183,7 +207,8 @@ async def handle_get_recent_tasks(websocket, redis_conn, num_tasks):
     for key in keys:
       content = await redis_conn.execute("lrange", key, "0", "-1")
       contents.append(json.loads(content[0].decode()))
-      timestamps += [timestamp for (timestamp, task, kind, info) in contents[-1] if task == "ray:task"]
+      timestamps += [timestamp for (timestamp, task, kind, info)
+                     in contents[-1] if task == "ray:task"]
 
     timestamps.sort()
     time_cutoff = timestamps[(-2 * num_tasks):][0]
@@ -197,36 +222,49 @@ async def handle_get_recent_tasks(websocket, redis_conn, num_tasks):
 
     num_tasks = 0
     task_data = [{"task_data": [],
-                  "num_workers": len(node_info[node_ip_address]["workers"])} for node_ip_address in node_ip_addresses]
+                  "num_workers": len(node_info[node_ip_address]["workers"])}
+                 for node_ip_address in node_ip_addresses]
     for i in range(len(keys)):
       worker_id, task_id = key_to_hex_identifiers(keys[i])
       data = contents[i]
       if worker_id not in worker_ids:
         # This case should be extremely rare.
-        raise Exception("A worker ID was not present in the list of worker IDs.")
-      node_ip_address = worker_info[worker_id][b"node_ip_address"].decode("ascii")
+        raise Exception("A worker ID was not present in the list of worker "
+                        "IDs.")
+      node_ip_address = (worker_info[worker_id][b"node_ip_address"]
+                         .decode("ascii"))
       worker_index = node_info[node_ip_address]["workers"].index(worker_id)
       node_index = node_ip_addresses.index(node_ip_address)
 
-      task_times = [timestamp for (timestamp, task, kind, info) in data if task == "ray:task"]
+      task_times = [timestamp for (timestamp, task, kind, info) in data
+                    if task == "ray:task"]
       if task_times[1] <= time_cutoff:
         continue
 
-      task_get_arguments_times = [timestamp for (timestamp, task, kind, info) in data if task == "ray:task:get_arguments"]
-      task_execute_times = [timestamp for (timestamp, task, kind, info) in data if task == "ray:task:execute"]
-      task_store_outputs_times = [timestamp for (timestamp, task, kind, info) in data if task == "ray:task:store_outputs"]
-      task_info = {"task": task_times,
-                   "get_arguments": task_get_arguments_times,
-                   "execute": task_execute_times,
-                   "store_outputs": task_store_outputs_times,
-                   "worker_index": worker_index,
-                   "node_ip_address": node_ip_address,
-                   "task_formatted_time": duration_to_string(task_times[1] - task_times[0]),
-                   "get_arguments_formatted_time": duration_to_string(task_get_arguments_times[1] - task_get_arguments_times[0])}
+      task_get_arguments_times = [timestamp for (timestamp, task, kind, info)
+                                  in data if task == "ray:task:get_arguments"]
+      task_execute_times = [timestamp for (timestamp, task, kind, info)
+                            in data if task == "ray:task:execute"]
+      task_store_outputs_times = [timestamp for (timestamp, task, kind, info)
+                                  in data if task == "ray:task:store_outputs"]
+      task_info = {
+          "task": task_times,
+          "get_arguments": task_get_arguments_times,
+          "execute": task_execute_times,
+          "store_outputs": task_store_outputs_times,
+          "worker_index": worker_index,
+          "node_ip_address": node_ip_address,
+          "task_formatted_time": duration_to_string(task_times[1] -
+                                                    task_times[0]),
+          "get_arguments_formatted_time":
+              duration_to_string(task_get_arguments_times[1] -
+                                 task_get_arguments_times[0])}
       if len(task_execute_times) == 2:
-        task_info["execute_formatted_time"] = duration_to_string(task_execute_times[1] - task_execute_times[0])
+        task_info["execute_formatted_time"] = duration_to_string(
+            task_execute_times[1] - task_execute_times[0])
       if len(task_store_outputs_times) == 2:
-        task_info["store_outputs_formatted_time"] = duration_to_string(task_store_outputs_times[1] - task_store_outputs_times[0])
+        task_info["store_outputs_formatted_time"] = duration_to_string(
+            task_store_outputs_times[1] - task_store_outputs_times[0])
       task_data[node_index]["task_data"].append(task_info)
       num_tasks += 1
     reply = {"min_time": min_time,
@@ -235,34 +273,41 @@ async def handle_get_recent_tasks(websocket, redis_conn, num_tasks):
              "task_data": task_data}
     await websocket.send(json.dumps(reply))
 
+
 async def send_heartbeat_payload(websocket):
   """Send heartbeat updates to the frontend every half second."""
   while True:
     reply = []
     for local_scheduler_id, local_scheduler in local_schedulers.items():
       current_time = time.time()
-      local_scheduler_info = {"local scheduler ID": local_scheduler_id,
-                              "time since heartbeat": duration_to_string(current_time - local_scheduler["last_heartbeat"]),
-                              "time since heartbeat numeric": str(current_time - local_scheduler["last_heartbeat"]),
-                              "node ip address": local_scheduler["node_ip_address"]}
+      local_scheduler_info = {
+          "local scheduler ID": local_scheduler_id,
+          "time since heartbeat":
+              (duration_to_string(current_time -
+               local_scheduler["last_heartbeat"])),
+          "time since heartbeat numeric":
+              str(current_time - local_scheduler["last_heartbeat"]),
+          "node ip address": local_scheduler["node_ip_address"]}
       reply.append(local_scheduler_info)
     # Send the payload to the frontend.
     await websocket.send(json.dumps(reply))
     # Wait for a little while so as not to overwhelm the frontend.
     await asyncio.sleep(0.5)
 
+
 async def send_heartbeats(websocket, redis_conn):
   # First update the local scheduler info locally.
   client_keys = await redis_conn.execute("keys", "CL:*")
-  clients = []
   for client_key in client_keys:
     client_fields = await hgetall_as_dict(redis_conn, client_key)
     if client_fields[b"client_type"] == b"local_scheduler":
       local_scheduler_id = hex_identifier(client_fields[b"ray_client_id"])
-      local_schedulers[local_scheduler_id] = {"node_ip_address": client_fields[b"node_ip_address"].decode("ascii"),
-                                              "local_scheduler_socket_name": client_fields[b"local_scheduler_socket_name"].decode("ascii"),
-                                              "aux_address": client_fields[b"aux_address"].decode("ascii"),
-                                              "last_heartbeat": -1 * np.inf}
+      local_schedulers[local_scheduler_id] = {
+          "node_ip_address": client_fields[b"node_ip_address"].decode("ascii"),
+          "local_scheduler_socket_name":
+              client_fields[b"local_scheduler_socket_name"].decode("ascii"),
+          "aux_address": client_fields[b"aux_address"].decode("ascii"),
+          "last_heartbeat": -1 * np.inf}
 
   # Subscribe to local scheduler heartbeats.
   await redis_conn.execute_pubsub("subscribe", "local_schedulers")
@@ -272,7 +317,8 @@ async def send_heartbeats(websocket, redis_conn):
 
   while True:
     msg = await redis_conn.pubsub_channels["local_schedulers"].get()
-    heartbeat = LocalSchedulerInfoMessage.GetRootAsLocalSchedulerInfoMessage(msg, 0)
+    heartbeat = LocalSchedulerInfoMessage.GetRootAsLocalSchedulerInfoMessage(
+        msg, 0)
     local_scheduler_id_bytes = heartbeat.DbClientId()
     local_scheduler_id = hex_identifier(local_scheduler_id_bytes)
     if local_scheduler_id not in local_schedulers:
@@ -281,6 +327,7 @@ async def send_heartbeats(websocket, redis_conn):
       continue
     local_schedulers[local_scheduler_id]["last_heartbeat"] = time.time()
 
+
 async def cache_data_from_redis(redis_ip_address, redis_port):
   """Open up ports to listen for new updates from Redis."""
   # TODO(richard): A lot of code needs to be ported in order to open new
@@ -288,6 +335,7 @@ async def cache_data_from_redis(redis_ip_address, redis_port):
 
   asyncio.ensure_future(listen_for_errors(redis_ip_address, redis_port))
 
+
 async def handle_get_log_files(websocket, redis_conn):
   reply = {}
   # First get all keys for the log file lists.
@@ -296,9 +344,11 @@ async def handle_get_log_files(websocket, redis_conn):
     node_ip_address = log_file_list_key.decode("ascii").split(":")[1]
     reply[node_ip_address] = {}
     # Get all of the log filenames for this node IP address.
-    log_filenames = await redis_conn.execute("lrange", log_file_list_key, 0, -1)
+    log_filenames = await redis_conn.execute("lrange", log_file_list_key, 0,
+                                             -1)
     for log_filename in log_filenames:
-      log_filename_key = "LOGFILE:{}:{}".format(node_ip_address, log_filename.decode("ascii"))
+      log_filename_key = "LOGFILE:{}:{}".format(node_ip_address,
+                                                log_filename.decode("ascii"))
       logfile = await redis_conn.execute("lrange", log_filename_key, 0, -1)
       logfile = [line.decode("ascii") for line in logfile]
       reply[node_ip_address][log_filename.decode("ascii")] = logfile
@@ -306,8 +356,10 @@ async def handle_get_log_files(websocket, redis_conn):
   # Send the reply back to the front end.
   await websocket.send(json.dumps(reply))
 
+
 async def serve_requests(websocket, path):
-  redis_conn = await aioredis.create_connection((redis_ip_address, redis_port), loop=loop)
+  redis_conn = await aioredis.create_connection((redis_ip_address, redis_port),
+                                                loop=loop)
   while True:
     command = json.loads(await websocket.recv())
     print("received command {}".format(command))
@@ -352,10 +404,10 @@ async def serve_requests(websocket, path):
                        "data_size": content[5].decode()})
       await websocket.send(json.dumps(result))
     elif command["command"] == "get-object-info":
-      # TODO(pcm): Get the object here (have to connect to ray) and ship content
-      # and type back to webclient. One challenge here is that the naive
-      # implementation will block the web ui backend, which is not ok if it is
-      # serving multiple users.
+      # TODO(pcm): Get the object here (have to connect to ray) and ship
+      # content and type back to webclient. One challenge here is that the
+      # naive implementation will block the web ui backend, which is not ok if
+      # it is serving multiple users.
       await websocket.send(json.dumps({"object_id": "none"}))
     elif command["command"] == "get-tasks":
       result = []
@@ -372,7 +424,8 @@ async def serve_requests(websocket, path):
         worker_id, task_id = key_to_hex_identifiers(key)
         content = await redis_conn.execute("lrange", key, "0", "-1")
         data = json.loads(content[0].decode())
-        begin_and_end_time = [timestamp for (timestamp, task, kind, info) in data if task == "ray:task"]
+        begin_and_end_time = [timestamp for (timestamp, task, kind, info)
+                              in data if task == "ray:task"]
         tasks[worker_id].append({"task_id": task_id,
                                  "start_task": min(begin_and_end_time),
                                  "end_task": max(begin_and_end_time)})
@@ -396,8 +449,8 @@ if __name__ == "__main__":
   redis_ip_address, redis_port = redis_address[0], int(redis_address[1])
 
   # The port here must match the value used by the frontend to connect over
-  # websockets. TODO(richard): Automatically increment the port if it is already
-  # taken.
+  # websockets. TODO(richard): Automatically increment the port if it is
+  # already taken.
   port = 8888
 
   loop.run_until_complete(cache_data_from_redis(redis_ip_address, redis_port))