cli: clean up outputs + prettify (#211)

Summary:
This cleans up the CLI output to make it be more readable and prettier by using colors.

* All stderr logs are prefixed with `torchx <timestamp> <log level>`. This allows distinguishing stdout with the local scheduler.
* runopts output has been reworked
* added `--log_level` CLI flag
* reduced log spam when using run

Pull Request resolved: #211

Test Plan:
NOTE: actual PR has "dim" timestamps not white/light gray since light gray does not render well at all on light backgrounds

### dim timestamps (actual)
![2021-09-29-183353_743x87_scrot](https://user-images.githubusercontent.com/909104/135371260-a35d3588-180d-4604-94e7-14466e044c6e.png)
![2021-09-29-183345_843x101_scrot](https://user-images.githubusercontent.com/909104/135371263-75c5489c-0fdb-4619-b34e-be4cdb5d3b74.png)

### run (wrong color timestamps)
![2021-09-29-181820_883x266_scrot](https://user-images.githubusercontent.com/909104/135369973-79d8c031-8f3f-41b7-b827-5e093cc471a2.png)

### log
![2021-09-29-181732_939x46_scrot](https://user-images.githubusercontent.com/909104/135369974-8a58486f-4e4f-42f4-8f5a-e24a26f3f3c4.png)

### runopts
![2021-09-29-181710_474x578_scrot](https://user-images.githubusercontent.com/909104/135369976-8f19293c-45b1-4297-b53a-96d7c218f7d8.png)

### log level 0  (wrong color timestamps)
![2021-09-29-182000_914x653_scrot](https://user-images.githubusercontent.com/909104/135369972-9ee2e536-1b5b-4650-ab6a-9a45a9ea74ca.png)

Reviewed By: kiukchung

Differential Revision: D31295428

Pulled By: d4l3k

fbshipit-source-id: a4eb6ffb0ff4744581e48559c8b61f9169877a19

Author: d4l3k

torchx/cli/cmd_log.py

@@ -17,20 +17,12 @@
 from pyre_extensions import none_throws
 from torchx import specs
 from torchx.cli.cmd_base import SubCommand
+from torchx.cli.colors import GREEN, ENDC
 from torchx.runner import Runner, get_runner
 from torchx.specs.api import make_app_handle
 
 logger: logging.Logger = logging.getLogger(__name__)
 
-# only print colors if outputting directly to a terminal
-if sys.stdout.isatty():
-    GREEN = "\033[32m"
-    ENDC = "\033[0m"
-else:
-    GREEN = ""
-    ENDC = ""
-
-
 ID_FORMAT = "SCHEDULER://[SESSION_NAME]/APP_ID/[ROLE_NAME/[REPLICA_IDS,...]]"
 
 

torchx/cli/cmd_run.py

@@ -158,9 +158,15 @@ def run(self, args: argparse.Namespace) -> None:
 
     def _wait_and_exit(self, runner: Runner, app_handle: str) -> None:
         logger.info("Waiting for the app to finish...")
+
         status = runner.wait(app_handle, wait_interval=1)
-        logger.info(status)
         if not status:
             raise RuntimeError(f"unknown status, wait returned {status}")
+
+        logger.info(f"Job finished: {status.state}")
+
         if status.state != specs.AppState.SUCCEEDED:
+            logger.error(status)
             sys.exit(1)
+        else:
+            logger.debug(status)

torchx/cli/cmd_runopts.py

@@ -9,6 +9,7 @@
 import logging
 
 from torchx.cli.cmd_base import SubCommand
+from torchx.cli.colors import GREEN, ENDC
 from torchx.runner.api import get_runner
 
 logger: logging.Logger = logging.getLogger(__name__)
@@ -24,11 +25,9 @@ def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
         )
 
     def run(self, args: argparse.Namespace) -> None:
-        scheduler = args.scheduler
+        filter = args.scheduler
         run_opts = get_runner().run_opts()
 
-        if not scheduler:
-            for scheduler, opts in run_opts.items():
-                logger.info(f"{scheduler}:\n{repr(opts)}")
-        else:
-            logger.info(repr(run_opts[scheduler]))
+        for scheduler, opts in run_opts.items():
+            if not filter or scheduler == filter:
+                print(f"{GREEN}{scheduler}{ENDC}:\n{repr(opts)}\n")

torchx/cli/colors.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import sys
+
+# only print colors if outputting directly to a terminal
+if sys.stdout.isatty():
+    GREEN = "\033[32m"
+    ORANGE = "\033[38:2:238:76:44m"
+    GRAY = "\033[2m"
+    ENDC = "\033[0m"
+else:
+    GREEN = ""
+    ORANGE = ""
+    GRAY = ""
+    ENDC = ""

torchx/cli/main.py

@@ -14,6 +14,7 @@
 from torchx.cli.cmd_run import CmdBuiltins, CmdRun
 from torchx.cli.cmd_runopts import CmdRunopts
 from torchx.cli.cmd_status import CmdStatus
+from torchx.cli.colors import ORANGE, GRAY, ENDC
 
 sub_parser_description = """Use the following commands to run operations, e.g.:
 torchx run ${JOB_NAME}
@@ -26,6 +27,12 @@ def create_parser() -> ArgumentParser:
     """
 
     parser = ArgumentParser(description="torchx CLI")
+    parser.add_argument(
+        "--log_level",
+        type=int,
+        help="Python logging log level",
+        default=logging.INFO,
+    )
     subparser = parser.add_subparsers(
         title="sub-commands",
         description=sub_parser_description,
@@ -49,13 +56,15 @@ def create_parser() -> ArgumentParser:
 
 
 def main(argv: List[str] = sys.argv[1:]) -> None:
+    parser = create_parser()
+    args = parser.parse_args(argv)
+
     logging.basicConfig(
-        level=logging.INFO,
-        format="%(message)s",
+        level=args.log_level,
+        format=f"{ORANGE}torchx{ENDC} {GRAY}%(asctime)s %(levelname)-8s{ENDC} %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
     )
 
-    parser = create_parser()
-    args = parser.parse_args(argv)
     if "func" not in args:
         parser.print_help()
         sys.exit(1)

torchx/schedulers/local_scheduler.py

@@ -231,7 +231,7 @@ def fetch(self, image: str) -> str:
         try:
             subprocess.run(["docker", "pull", image], check=True)
         except Exception as e:
-            print(f"failed to fetch image {image}, falling back to local: {e}")
+            log.warning(f"failed to fetch image {image}, falling back to local: {e}")
         return ""
 
     def get_replica_param(
@@ -438,7 +438,7 @@ def _fmt_io_filename(std_io: Optional[TextIO]) -> str:
         with open(os.path.join(self.log_dir, "SUCCESS"), "w") as fp:
             fp.write(info_str)
 
-        log.info(f"Successfully closed app_id: {self.id}.\n{info_str}")
+        log.debug(f"Successfully closed app_id: {self.id}.\n{info_str}")
 
     def __repr__(self) -> str:
         role_to_pid = {}
@@ -639,7 +639,7 @@ def _popen(
             env["PATH"] = os.pathsep.join([p for p in path if p])  # remove empty str
 
         args_pfmt = pprint.pformat(asdict(replica_params), indent=2, width=80)
-        log.info(f"Running {role_name} (replica {replica_id}):\n {args_pfmt}")
+        log.debug(f"Running {role_name} (replica {replica_id}):\n {args_pfmt}")
 
         proc = subprocess.Popen(
             args=replica_params.args,

torchx/specs/api.py

@@ -657,25 +657,32 @@ def resolve(self, config: RunConfig) -> RunConfig:
         return resolved_cfg
 
     def __repr__(self) -> str:
-        # make it a pretty printable dict
-        pretty_opts = {}
-        for cfg_key, runopt in self._opts.items():
-            key = f"*{cfg_key}" if runopt.is_required else cfg_key
-            opt = {"type": get_type_name(runopt.opt_type)}
-            if runopt.is_required:
-                opt["required"] = "True"
-            else:
-                opt["default"] = str(runopt.default)
-            opt["help"] = runopt.help
-
-            pretty_opts[key] = opt
-        import pprint
-
-        return pprint.pformat(
-            pretty_opts,
-            indent=2,
-            width=80,
-        )
+        required = [(key, opt) for key, opt in self._opts.items() if opt.is_required]
+        optional = [
+            (key, opt) for key, opt in self._opts.items() if not opt.is_required
+        ]
+
+        out = "    usage:\n        "
+        for i, (key, opt) in enumerate(required + optional):
+            contents = f"{key}={key.upper()}"
+            if not opt.is_required:
+                contents = f"[{contents}]"
+            if i > 0:
+                contents = "," + contents
+            out += contents
+
+        sections = [("required", required), ("optional", optional)]
+
+        for section, opts in sections:
+            if len(opts) == 0:
+                continue
+            out += f"\n\n    {section} arguments:"
+            for key, opt in opts:
+                default = "" if opt.is_required else f", {opt.default}"
+                out += f"\n        {key}={key.upper()} ({get_type_name(opt.opt_type)}{default})"
+                out += f"\n            {opt.help}"
+
+        return out
 
 
 class InvalidRunConfigException(Exception):