[core][state][no_early_kickoff] Add "humanify" feature to StateSchema

python/ray/experimental/state/common.py

@@ -1,3 +1,5 @@
+import dataclasses
+import datetime
 import json
 import logging
 import sys
@@ -156,14 +158,15 @@ def state_column(*, filterable: bool, detail: bool = False, **kwargs):
         kwargs: The same kwargs for the `dataclasses.field` function.
     """
     m = {"detail": detail, "filterable": filterable}
-
     # Default for detail field is None since it could be missing.
     if detail and "default" not in kwargs:
         kwargs["default"] = None
 
     if "metadata" in kwargs:
+        # Metadata explicitly specified, so add detail and filterable if missing.
         kwargs["metadata"].update(m)
     else:
+        # Metadata not explicitly specified, so add it.
         kwargs["metadata"] = m
     return field(**kwargs)
 
@@ -193,8 +196,36 @@ class State(StateSchema):
     # Returns {"column_a", "column_b"}
     s.columns()
     ```
+
+    In addition, the schema also provides a humanify abstract method to
+    convert the state object into something human readable, ready for printing.
+
+    Subclasses should override this method, providing logic to convert its own fields
+    to something human readable, packaged and returned in a dict.
+
+    Each field that wants to be humanified should include a 'format_fn' key in its
+    metadata dictionary.
     """
 
+    @classmethod
+    def humanify(cls, state: dict) -> dict:
+        """Convert the given state object into something human readable."""
+        if not cls.is_valid_state(state):
+            return state  # return the original state
+        for f in fields(cls):
+            if f.metadata.get("format_fn") is not None and state[f.name] is not None:
+                state[f.name] = f.metadata["format_fn"](state[f.name])
+        return state
+
+    @classmethod
+    def is_valid_state(cls, state: dict) -> bool:
+        """Checks to make sure state has all keys in the current schema class."""
+        cols = cls.list_columns(detail=True)
+        for col in cols:
+            if col not in state:
+                return False
+        return True
+
     @classmethod
     def list_columns(cls, detail: bool = True) -> List[str]:
         """Return a list of columns."""
@@ -516,6 +547,35 @@ class ClusterEventState(StateSchema):
     custom_fields: Optional[dict] = state_column(filterable=False, detail=True)
 
 
+class Humanify:
+    """A class containing default methods to convert units into a human readable string."""
+
+    def timestamp(x: float):
+        """Converts miliseconds to a datetime object."""
+        return str(datetime.datetime.fromtimestamp(x / 1000))
+
+    def memory(x: int):
+        """Converts raw bytes to a human readable memory size."""
+        if x >= 2**30:
+            return str(x / (2**30)) + " GiB"
+        elif x >= 2**20:
+            return str(x / (2**20)) + " MiB"
+        elif x >= 2**10:
+            return str(x / (2**10)) + " KiB"
+        return str(x) + " B"
+
+    def duration(x: int):
+        """Converts miliseconds to a human readable duration."""
+        return str(datetime.timedelta(milliseconds=x))
+
+    def events(events: List[dict]):
+        """Converts a list of task events into a human readable format."""
+        for event in events:
+            if "created_ms" in event:
+                event["created_ms"] = Humanify.timestamp(event["created_ms"])
+        return events
+
+
 @dataclass(init=True)
 class TaskState(StateSchema):
     """Task State"""
@@ -571,15 +631,29 @@ class TaskState(StateSchema):
     #: The list of events of the given task.
     #: Refer to src/ray/protobuf/common.proto for a detailed explanation of the state
     #: breakdowns and typical state transition flow.
-    events: Optional[List[dict]] = state_column(detail=True, filterable=False)
+    events: Optional[List[dict]] = state_column(
+        detail=True, filterable=False, metadata={"format_fn": Humanify.events}
+    )
     #: The list of profile events of the given task.
     profiling_data: Optional[dict] = state_column(detail=True, filterable=False)
     #: The time when the task is created. A Unix timestamp in ms.
-    creation_time_ms: Optional[int] = state_column(detail=True, filterable=False)
+    creation_time_ms: Optional[int] = state_column(
+        detail=True,
+        filterable=False,
+        metadata={"format_fn": Humanify.timestamp},
+    )
     #: The time when the task starts to run. A Unix timestamp in ms.
-    start_time_ms: Optional[int] = state_column(detail=True, filterable=False)
+    start_time_ms: Optional[int] = state_column(
+        detail=True,
+        filterable=False,
+        metadata={"format_fn": Humanify.timestamp},
+    )
     #: The time when the task is finished or failed. A Unix timestamp in ms.
-    end_time_ms: Optional[int] = state_column(detail=True, filterable=False)
+    end_time_ms: Optional[int] = state_column(
+        detail=True,
+        filterable=False,
+        metadata={"format_fn": Humanify.timestamp},
+    )
     #: The task logs info, e.g. offset into the worker log file when the task
     #: starts/finishes.
     task_log_info: Optional[dict] = state_column(detail=True, filterable=False)

python/ray/experimental/state/state_cli.py

@@ -172,6 +172,9 @@ def output_with_format(
     format: AvailableFormat = AvailableFormat.DEFAULT,
     detail: bool = False,
 ) -> str:
+    # humanify all input state data
+    state_data = [schema.humanify(state) for state in state_data]
+
     if format == AvailableFormat.DEFAULT:
         return get_table_output(state_data, schema, detail)
     if format == AvailableFormat.YAML:
@@ -292,8 +295,10 @@ def format_get_api_output(
 ) -> str:
     if not state_data or isinstance(state_data, list) and len(state_data) == 0:
         return f"Resource with id={id} not found in the cluster."
+
     if not isinstance(state_data, list):
         state_data = [state_data]
+
     state_data = [dataclasses.asdict(state) for state in state_data]
     return output_with_format(state_data, schema=schema, format=format, detail=True)
 

python/ray/tests/test_state_api.py

@@ -1,4 +1,6 @@
+import datetime
 import os
+import re
 import time
 import json
 import sys
@@ -8,6 +10,7 @@
 from unittest.mock import MagicMock
 
 import pytest
+from ray.experimental.state.common import Humanify
 from ray._private.gcs_utils import GcsAioClient
 import yaml
 from click.testing import CliRunner
@@ -1696,6 +1699,18 @@ def ready(self):
     assert result.total == 6
 
 
+def test_humanify():
+    raw_bytes = 1024
+    assert Humanify.memory(raw_bytes) == "1.0 KiB"
+    raw_bytes *= 1024
+    assert Humanify.memory(raw_bytes) == "1.0 MiB"
+    raw_bytes *= 1024
+    assert Humanify.memory(raw_bytes) == "1.0 GiB"
+    timestamp = 1610000000
+    assert Humanify.timestamp(timestamp) == "1970-01-19 07:13:20"
+    assert Humanify.duration(timestamp) == "18 days, 15:13:20"
+
+
 @pytest.mark.asyncio
 async def test_state_data_source_client_limit_distributed_sources(ray_start_cluster):
     cluster = ray_start_cluster