schedulers,cli: persist newline breaks in log_iter

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/schedulers/api.py

@@ -302,3 +302,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/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 (
@@ -394,7 +395,7 @@ def log_iter(
             if len(logs) == 0:
                 logs = []
             else:
-                logs = logs.split("\n")
+                logs = split_lines(logs)
 
         logs = map(_to_str, logs)
 
@@ -407,8 +408,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 (
@@ -599,7 +600,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"])

torchx/schedulers/test/docker_scheduler_test.py

@@ -201,8 +201,8 @@ def test_docker_logs(self) -> None:
             self.assertEqual(
                 logs,
                 [
-                    "foo",
-                    "bar",
+                    "foo\n",
+                    "bar\n",
                 ],
             )
             logs = list(
@@ -216,7 +216,7 @@ def test_docker_logs(self) -> None:
             self.assertEqual(
                 logs,
                 [
-                    "bar",
+                    "bar\n",
                 ],
             )
 
@@ -249,8 +249,8 @@ def test_docker_logs(self) -> None:
             self.assertEqual(
                 logs,
                 [
-                    "foo",
-                    "bar",
+                    "foo\n",
+                    "bar\n",
                 ],
             )
 
@@ -268,8 +268,8 @@ def test_docker_logs_streams(self) -> None:
             self.assertEqual(
                 logs,
                 {
-                    "stdout",
-                    "stderr",
+                    "stdout\n",
+                    "stderr\n",
                 },
             )
 
@@ -281,8 +281,8 @@ def test_docker_logs_streams(self) -> None:
             self.assertEqual(
                 logs,
                 {
-                    "stdout",
-                    "stderr",
+                    "stdout\n",
+                    "stderr\n",
                 },
             )
 
@@ -294,7 +294,7 @@ def test_docker_logs_streams(self) -> None:
             self.assertEqual(
                 logs,
                 [
-                    "stderr",
+                    "stderr\n",
                 ],
             )
 
@@ -306,7 +306,7 @@ def test_docker_logs_streams(self) -> None:
             self.assertEqual(
                 logs,
                 [
-                    "stdout",
+                    "stdout\n",
                 ],
             )
 

torchx/schedulers/test/kubernetes_scheduler_test.py

@@ -511,8 +511,8 @@ def test_log_iter(self, read_namespaced_pod_log: MagicMock) -> None:
         self.assertEqual(
             list(lines),
             [
-                "foo reg",
-                "bar reg",
+                "foo reg\n",
+                "bar reg\n",
             ],
         )
         call = read_namespaced_pod_log.call_args

torchx/schedulers/test/local_scheduler_test.py

@@ -325,7 +325,7 @@ def test_submit_inherit_parent_envs(self) -> None:
         app = AppDef(name="check_foo_env_var", roles=[role])
         app_id = self.scheduler.submit(app, {"log_dir": self.test_dir})
         for line in self.scheduler.log_iter(app_id, "echo_foo"):
-            self.assertEqual("bar", line)
+            self.assertEqual("bar\n", line)
 
         desc = self.wait(app_id, self.scheduler)
         assert desc is not None
@@ -429,7 +429,7 @@ def test_submit_override_parent_env(self) -> None:
         app = AppDef(name="check_foo_env_var", roles=[role])
         app_id = self.scheduler.submit(app, {"log_dir": self.test_dir})
         for line in self.scheduler.log_iter(app_id, "echo_foo"):
-            self.assertEqual("new_bar", line)
+            self.assertEqual("new_bar\n", line)
 
         desc = self.wait(app_id, self.scheduler)
         assert desc is not None
@@ -598,20 +598,20 @@ def test_log_iterator(self) -> None:
         app_id = self.scheduler.submit(app, cfg)
 
         for i, line in enumerate(self.scheduler.log_iter(app_id, "role1", k=0)):
-            self.assertEqual(str(i), line)
+            self.assertEqual(str(i), line.strip())
 
         # since and until ignored
         for i, line in enumerate(
             self.scheduler.log_iter(
                 app_id, "role1", k=0, since=datetime.now(), until=datetime.now()
             )
         ):
-            self.assertEqual(str(i), line)
+            self.assertEqual(str(i), line.strip())
 
         for i, line in enumerate(
             self.scheduler.log_iter(app_id, "role1", k=0, regex=r"[02468]")
         ):
-            self.assertEqual(str(i * 2), line)
+            self.assertEqual(str(i * 2), line.strip())
 
     def test_log_iterator_no_log_dir(self) -> None:
         role = Role(

torchx/schedulers/test/slurm_scheduler_test.py

@@ -373,7 +373,7 @@ def test_log_iter(self, run: MagicMock) -> None:
                         since=datetime.datetime.now(),
                     )
                 )
-                self.assertEqual(logs, ["hello", "world"])
+                self.assertEqual(logs, ["hello\n", "world\n"])
 
                 with open(os.path.join(job_dir, "slurm-54-echo-1.err"), "wt") as f:
                     f.write("foo\nbar\n")
@@ -387,7 +387,7 @@ def test_log_iter(self, run: MagicMock) -> None:
                     )
                 )
 
-                self.assertEqual(logs, ["foo", "bar"])
+                self.assertEqual(logs, ["foo\n", "bar\n"])
 
                 # no stream specified should default to STDERR
                 logs = list(
@@ -397,7 +397,7 @@ def test_log_iter(self, run: MagicMock) -> None:
                         1,
                     )
                 )
-                self.assertEqual(logs, ["foo", "bar"])
+                self.assertEqual(logs, ["foo\n", "bar\n"])
 
         with self.assertRaises(ValueError):
             scheduler.log_iter("54", "echo", 1, streams=Stream.COMBINED)