support AppStatus json output format

Summary:
Enable json output format to show AppStatus

What added:
1. Add `--json` subcommand for `status`, default is False
2. Parse both AppStatus and RoleStatus to dictionary format
3. json.dump dictionary in output
4. Add unit test

Differential Revision: D73180040

Author: nghuiqin

torchx/cli/cmd_status.py

@@ -46,6 +46,11 @@ def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
         subparser.add_argument(
             "--roles", type=str, default="", help="comma separated roles to filter"
         )
+        subparser.add_argument(
+            "--json",
+            action="store_true",
+            help="output the status in JSON format",
+        )
 
     def run(self, args: argparse.Namespace) -> None:
         app_handle = args.app_handle
@@ -54,7 +59,7 @@ def run(self, args: argparse.Namespace) -> None:
         app_status = runner.status(app_handle)
         filter_roles = parse_list_arg(args.roles)
         if app_status:
-            print(app_status.format(filter_roles))
+            print(app_status.format(filter_roles, args.json))
         else:
             logger.error(
                 f"AppDef: {app_id},"

torchx/specs/api.py

@@ -538,6 +538,12 @@ class RoleStatus:
     role: str
     replicas: List[ReplicaStatus]
 
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "role": self.role,
+            "replicas": [asdict(replica) for replica in self.replicas],
+        }
+
 
 @dataclass
 class AppStatus:
@@ -657,9 +663,20 @@ def _format_role_status(
             replica_data += self._format_replica_status(replica)
         return f"{replica_data}"
 
+    def _to_json(self, roles: List[RoleStatus]) -> Dict[str, Any]:
+        return {
+            "state": str(self.state),
+            "num_restarts": self.num_restarts,
+            "roles": [role_status.to_dict() for role_status in roles],
+            "msg": self.msg,
+            "structured_error_msg": self.structured_error_msg,
+            "url": self.ui_url,
+        }
+
     def format(
         self,
         filter_roles: Optional[List[str]] = None,
+        as_json: bool = False,
     ) -> str:
         """
         Format logs for app status. The app status include:
@@ -669,19 +686,27 @@ def format(
             4. Msg: Arbitrary text message the scheduler returned.
             5. Structured Error Msg: Json response error msg.
             6. UI URL: Application URL
+
+        Args:
+            filter_roles: Optional list of role names to filter by
+            as_json: Show the output in json format
         """
         roles_data = ""
         roles = self._get_role_statuses(self.roles, filter_roles)
-        for role_status in roles:
-            roles_data += self._format_role_status(role_status)
-        return Template(_APP_STATUS_FORMAT_TEMPLATE).substitute(
-            state=self.state,
-            num_restarts=self.num_restarts,
-            roles=roles_data,
-            msg=self.msg,
-            structured_error_msg=self.structured_error_msg,
-            url=self.ui_url,
-        )
+
+        if as_json:
+            return json.dumps(self._to_json(roles), indent=4)
+        else:
+            for role_status in roles:
+                roles_data += self._format_role_status(role_status)
+            return Template(_APP_STATUS_FORMAT_TEMPLATE).substitute(
+                state=self.state,
+                num_restarts=self.num_restarts,
+                roles=roles_data,
+                msg=self.msg,
+                structured_error_msg=self.structured_error_msg,
+                url=self.ui_url,
+            )
 
 
 class AppStatusError(Exception):

torchx/specs/test/api_test.py

@@ -176,6 +176,44 @@ def test_format_app_status(self) -> None:
         # Split and compare to aviod AssertionError.
         self.assertEqual(expected_message.split(), actual_message.split())
 
+    def test_format_app_status_in_json(self) -> None:
+        app_status = self._get_test_app_status()
+        actual_message = app_status.format(as_json=True)
+        error_msg = '{\\"message\\":{\\"message\\":\\"error\\",\\"errorCode\\":-1,\\"extraInfo\\":{\\"timestamp\\":1293182}}}'
+        expected_message = (
+            """{
+    "state": "RUNNING",
+    "num_restarts": 0,
+    "roles": [
+        {
+            "role": "worker",
+            "replicas": [
+                {
+                    "id": 0,
+                    "state": 5,
+                    "role": "worker",
+                    "hostname": "localhost",
+                    "structured_error_msg": "%s"
+                },
+                {
+                    "id": 1,
+                    "state": 3,
+                    "role": "worker",
+                    "hostname": "localhost",
+                    "structured_error_msg": "<NONE>"
+                }
+            ]
+        }
+    ],
+    "msg": "",
+    "structured_error_msg": "<NONE>",
+    "url": null
+}"""
+            % error_msg
+        )
+        print(expected_message)
+        self.assertEqual(expected_message, actual_message)
+
 
 class ResourceTest(unittest.TestCase):
     def test_copy_resource(self) -> None: