schedulers,cli: persist newline breaks in log_iter (#425)

Summary:
This resolves #424

This makes it so the torchx scheduler `log_iter` method keeps the line breaks so downstream log streams can handle them gracefully. The current solution strips all `\n` characters and always adds them so it makes it impossible to do streaming visualizations of progress bars which use `\r` without a new line break.

WARNING: This is a change in the log_iter interface and all schedulers/downstream consumers will need to be updated.

If someone is logging from multiple workers this gets dangerous since the progress bar `\r` lines can clobber each other.

Pull Request resolved: #425

Test Plan:
(torchx-3.10.2) tristanr@tristanr-arch2 ~/D/torchx-proj> torchx run --scheduler local_docker --wait --log utils.python --script test_tqdm.py
torchx 2022-03-15 14:26:42 INFO     loaded configs from /home/tristanr/Developer/torchx-proj/.torchxconfig
torchx 2022-03-15 14:26:42 INFO     Building workspace: file:///home/tristanr/Developer/torchx-proj for role[0]: python, image: ghcr.io/pytorch/torchx:0.1.2
dev0
torchx 2022-03-15 14:26:43 INFO     Done building workspace
torchx 2022-03-15 14:26:43 INFO     New image: sha256:9cfaf70f7143b4caef383b46c23635eaf001cbd3d9ff55335aa1ff8c5e236388 built from workspace
local_docker://torchx/torchx_utils_python-bprr9rb4k764nd
torchx 2022-03-15 14:26:44 INFO     Waiting for the app to finish...
python/0 100%|██████████| 100/100 [00:03<00:00, 32.95it/s]
torchx 2022-03-15 14:26:48 INFO     Job finished: SUCCEEDED
(torchx-3.10.2) tristanr@tristanr-arch2 ~/D/torchx-proj> torchx run --scheduler local_cwd --wait --log utils.python --script test_tqdm.py
torchx 2022-03-15 14:26:52 INFO     loaded configs from /home/tristanr/Developer/torchx-proj/.torchxconfig
torchx 2022-03-15 14:26:52 INFO     Log files located in: /tmp/torchx_0nqvqm1d/torchx/torchx_utils_python-x217jjqhbkkrgd/python/0
local_cwd://torchx/torchx_utils_python-x217jjqhbkkrgd
torchx 2022-03-15 14:26:52 INFO     Waiting for the app to finish...
python/0 100%|██████████| 100/100 [00:03<00:00, 32.95it/s]
torchx 2022-03-15 14:26:56 INFO     Job finished: SUCCEEDED

Reviewed By: kiukchung

Differential Revision: D34907682

Pulled By: d4l3k

fbshipit-source-id: 1750a301ca7c8319df4fbeca35fa8eb29b8ecaae

Author: d4l3k

torchx/cli/cmd_log.py

@@ -35,6 +35,17 @@ def validate(job_identifier: str) -> None:
         sys.exit(1)
 
 
+def _prefix_line(prefix: str, line: str) -> str:
+    """
+    _prefix_line ensure the prefix is still present even when dealing with return characters
+    """
+    if "\r" in line:
+        line = line.replace("\r", f"\r{prefix}")
+    if not line.startswith("\r"):
+        line = f"{prefix}{line}"
+    return line
+
+
 def print_log_lines(
     file: TextIO,
     runner: Runner,
@@ -55,7 +66,8 @@ def print_log_lines(
             should_tail=should_tail,
             streams=streams,
         ):
-            print(f"{GREEN}{role_name}/{replica_id}{ENDC} {line}", file=file)
+            prefix = f"{GREEN}{role_name}/{replica_id}{ENDC} "
+            print(_prefix_line(prefix, line), file=file, end="")
     except Exception as e:
         exceptions.put(e)
         raise

torchx/cli/test/cmd_log_test.py

@@ -65,7 +65,7 @@ def log_lines(
         if regex is None:
             regex = ".*"
 
-        log_lines = ["INFO foo", "ERROR bar", "WARN baz"]
+        log_lines = ["INFO foo\n", "ERROR bar\n", "WARN baz\n"]
         return iter([line for line in log_lines if re.match(regex, line)])
 
 

torchx/runner/api.py

@@ -469,28 +469,36 @@ def log_lines(
                      if the scheduler has already totally or partially purged log records
                      for the application.
 
+        Return lines will include whitespace characters such as ``\\n`` or
+        ``\\r``. When outputting the lines you should make sure to avoid adding
+        extra newline characters.
+
         Usage:
 
-        ::
+        .. code:: python
 
-         app_handle = session.run(app, scheduler="local", cfg=Dict[str, ConfigValue]())
+            app_handle = session.run(app, scheduler="local", cfg=Dict[str, ConfigValue]())
 
-         print("== trainer node 0 logs ==")
-         for line in session.log_lines(app_handle, "trainer", k=0):
-            print(line)
+            print("== trainer node 0 logs ==")
+            for line in session.log_lines(app_handle, "trainer", k=0):
+               # for prints newlines will already be present in the line
+               print(line, end="")
+
+               # when writing to a file nothing extra is necessary
+               f.write(line)
 
         Discouraged anti-pattern:
 
-        ::
+        .. code:: python
 
-         # DO NOT DO THIS!
-         # parses accuracy metric from log and reports it for this experiment run
-         accuracy = -1
-         for line in session.log_lines(app_handle, "trainer", k=0):
-            if matches_regex(line, "final model_accuracy:[0-9]*"):
-                accuracy = parse_accuracy(line)
-                break
-         report(experiment_name, accuracy)
+            # DO NOT DO THIS!
+            # parses accuracy metric from log and reports it for this experiment run
+            accuracy = -1
+            for line in session.log_lines(app_handle, "trainer", k=0):
+               if matches_regex(line, "final model_accuracy:[0-9]*"):
+                   accuracy = parse_accuracy(line)
+                   break
+            report(experiment_name, accuracy)
 
         Args:
             app_handle: application handle

torchx/schedulers/api.py

@@ -263,6 +263,11 @@ def log_iter(
         7. Some schedulers may support line cursors by supporting ``__getitem__``
            (e.g. ``iter[50]`` seeks to the 50th log line).
 
+        8. Whitespace is preserved, each new line should include ``\n``. To
+            support interactive progress bars the returned lines don't need to
+            include ``\n`` but should then be printed without a newline to
+            correctly handle ``\r`` carriage returns.
+
         Args:
             streams: The IO output streams to select.
                 One of: combined, stdout, stderr.
@@ -302,3 +307,19 @@ def filter_regex(regex: str, data: Iterable[str]) -> Iterable[str]:
 
     r = re.compile(regex)
     return filter(lambda datum: r.search(datum), data)
+
+
+def split_lines(text: str) -> List[str]:
+    """
+    split_lines splits the string by new lines and keeps the new line characters.
+    """
+    lines = []
+    while len(text) > 0:
+        idx = text.find("\n")
+        if idx >= 0:
+            lines.append(text[: idx + 1])
+            text = text[idx + 1 :]
+        else:
+            lines.append(text)
+            break
+    return lines

torchx/schedulers/aws_batch_scheduler.py

@@ -517,7 +517,7 @@ def _stream_events(
             next_token = response["nextForwardToken"]
 
             for event in response["events"]:
-                yield event["message"]
+                yield event["message"] + "\n"
 
 
 def create_scheduler(session_name: str, **kwargs: object) -> AWSBatchScheduler:

torchx/schedulers/docker_scheduler.py

@@ -20,6 +20,7 @@
     Scheduler,
     Stream,
     filter_regex,
+    split_lines,
 )
 from torchx.schedulers.ids import make_unique
 from torchx.specs.api import (
@@ -425,7 +426,7 @@ def log_iter(
             if len(logs) == 0:
                 logs = []
             else:
-                logs = logs.split("\n")
+                logs = split_lines(logs)
 
         logs = map(_to_str, logs)
 
@@ -438,8 +439,6 @@ def log_iter(
 def _to_str(a: Union[str, bytes]) -> str:
     if isinstance(a, bytes):
         a = a.decode("utf-8")
-    if a.endswith("\n"):
-        a = a[:-1]
     return a
 
 

torchx/schedulers/kubernetes_scheduler.py

@@ -49,6 +49,7 @@
     Scheduler,
     Stream,
     filter_regex,
+    split_lines,
 )
 from torchx.schedulers.ids import make_unique
 from torchx.specs.api import (
@@ -640,7 +641,7 @@ def log_iter(
             iterator = w.stream(core_api.read_namespaced_pod_log, **args)
         else:
             resp = core_api.read_namespaced_pod_log(**args)
-            iterator = resp.strip().split("\n")
+            iterator = split_lines(resp)
 
         if regex:
             return filter_regex(regex, iterator)

torchx/schedulers/local_scheduler.py

@@ -1020,7 +1020,7 @@ def __iter__(self) -> "LogIterator":
             self._check_finished()  # check to see if app has finished running
 
             if os.path.isfile(self._log_file):
-                self._log_fp = open(self._log_file, "r")  # noqa: P201
+                self._log_fp = open(self._log_file, "rt", newline="\n")  # noqa: P201
                 break
 
             if self._app_finished:
@@ -1048,7 +1048,6 @@ def __next__(self) -> str:
                 time.sleep(0.1)
                 self._check_finished()
             else:
-                line = line.rstrip("\n")  # strip the trailing newline
                 if re.match(self._regex, line):
                     return line
 

torchx/schedulers/ray_scheduler.py

@@ -21,6 +21,7 @@
     DescribeAppResponse,
     Scheduler,
     Stream,
+    split_lines,
 )
 from torchx.schedulers.ids import make_unique
 from torchx.schedulers.ray.ray_common import RayActor
@@ -350,7 +351,7 @@ def log_iter(
             addr, app_id = app_id.split("-")
             client: JobSubmissionClient = JobSubmissionClient(f"http://{addr}")
             logs: str = client.get_job_logs(app_id)
-            return logs.split("\n")
+            return split_lines(logs)
 
     def create_scheduler(session_name: str, **kwargs: Any) -> RayScheduler:
         if not has_ray():  # pragma: no cover

torchx/schedulers/test/api_test.py

@@ -11,7 +11,12 @@
 from typing import Iterable, Mapping, Optional, Union
 from unittest.mock import MagicMock, patch
 
-from torchx.schedulers.api import DescribeAppResponse, Scheduler, Stream
+from torchx.schedulers.api import (
+    DescribeAppResponse,
+    Scheduler,
+    Stream,
+    split_lines,
+)
 from torchx.specs.api import (
     NULL_RESOURCE,
     AppDef,
@@ -152,3 +157,9 @@ def test_close_twice(self) -> None:
         scheduler_mock.close()
         scheduler_mock.close()
         # nothing to validate explicitly, just that no errors are raised
+
+    def test_split_lines(self) -> None:
+        self.assertEqual(split_lines(""), [])
+        self.assertEqual(split_lines("\n"), ["\n"])
+        self.assertEqual(split_lines("foo\nbar"), ["foo\n", "bar"])
+        self.assertEqual(split_lines("foo\nbar\n"), ["foo\n", "bar\n"])