Rebase on main @5d35e4a3 (#696)

tests/test_tracking_cli.py

@@ -62,7 +62,7 @@ def test_taskcluster(wandb_mock, getargs_mock, caplog, samples_dir, tmp_dir):
     assert [(level, message) for _module, level, message in caplog.record_tuples] == [
         (logging.INFO, "Reading logs stream."),
         (logging.INFO, "Detected Marian version 1.10"),
-        (logging.INFO, "Successfully parsed 588 lines"),
+        (logging.INFO, "Successfully parsed 1528 lines"),
         (logging.INFO, "Found 102 training entries"),
         (logging.INFO, "Found 34 validation entries"),
     ]
@@ -96,7 +96,7 @@ def test_experiments_marian_1_10(wandb_mock, getargs_mock, caplog, samples_dir,
             ),
             (logging.INFO, "Reading logs stream."),
             (logging.INFO, "Detected Marian version 1.10"),
-            (logging.INFO, "Successfully parsed 1002 lines"),
+            (logging.INFO, "Successfully parsed 1878 lines"),
             (logging.INFO, "Found 550 training entries"),
             (logging.INFO, "Found 108 validation entries"),
             # teacher-finetuned0
@@ -105,7 +105,7 @@ def test_experiments_marian_1_10(wandb_mock, getargs_mock, caplog, samples_dir,
                 f"Parsing folder {samples_dir}/experiments_1_10/models/en-nl/prod/teacher-finetuned0",
             ),
             (logging.INFO, "Reading logs stream."),
-            (logging.INFO, "Successfully parsed 993 lines"),
+            (logging.INFO, "Successfully parsed 1944 lines"),
             (logging.INFO, "Found 567 training entries"),
             (logging.INFO, "Found 189 validation entries"),
             # teacher-finetuned1
@@ -114,7 +114,7 @@ def test_experiments_marian_1_10(wandb_mock, getargs_mock, caplog, samples_dir,
                 f"Parsing folder {samples_dir}/experiments_1_10/models/en-nl/prod/teacher-finetuned1",
             ),
             (logging.INFO, "Reading logs stream."),
-            (logging.INFO, "Successfully parsed 1000 lines"),
+            (logging.INFO, "Successfully parsed 1963 lines"),
             (logging.INFO, "Found 573 training entries"),
             (logging.INFO, "Found 191 validation entries"),
             # Publish group files and quantized/evaluated metrics
@@ -189,15 +189,15 @@ def test_experiments_marian_1_12(wandb_mock, getargs_mock, caplog, samples_dir,
                 f"Parsing folder {samples_dir}/experiments_1_12/models/fi-en/opusprod/student",
             ),
             (logging.INFO, "Reading logs stream."),
-            (logging.INFO, "Successfully parsed 786 lines"),
+            (logging.INFO, "Successfully parsed 1533 lines"),
             (logging.INFO, "Found 405 training entries"),
             (logging.INFO, "Found 79 validation entries"),
             (
                 logging.INFO,
                 f"Parsing folder {samples_dir}/experiments_1_12/models/fi-en/opusprod/student-finetuned",
             ),
             (logging.INFO, "Reading logs stream."),
-            (logging.INFO, "Successfully parsed 660 lines"),
+            (logging.INFO, "Successfully parsed 1174 lines"),
             (logging.INFO, "Found 330 training entries"),
             (logging.INFO, "Found 64 validation entries"),
             (
@@ -275,7 +275,7 @@ def test_taskcluster_wandb_initialization_failure(
             logging.ERROR,
             "WandB client could not be initialized: Invalid credentials. No data will be published.",
         ),
-        (logging.INFO, "Successfully parsed 588 lines"),
+        (logging.INFO, "Successfully parsed 1528 lines"),
         (logging.INFO, "Found 102 training entries"),
         (logging.INFO, "Found 34 validation entries"),
     ]
@@ -320,7 +320,7 @@ def test_taskcluster_wandb_log_failures(wandb_mock, getargs_mock, caplog, sample
         (logging.ERROR, "Error publishing training epoch using WandB: Unexpected failure"),
         (logging.ERROR, "Error publishing validation epoch using WandB: Unexpected failure"),
     ] * 34 + [
-        (logging.INFO, "Successfully parsed 588 lines"),
+        (logging.INFO, "Successfully parsed 1528 lines"),
         (logging.INFO, "Found 102 training entries"),
         (logging.INFO, "Found 34 validation entries"),
     ]
@@ -359,7 +359,7 @@ def test_taskcluster_wandb_disabled(wandb_mock, getargs_mock, caplog, samples_di
         ),
         (logging.INFO, "Reading logs stream."),
         (logging.INFO, "Detected Marian version 1.10"),
-        (logging.INFO, "Successfully parsed 588 lines"),
+        (logging.INFO, "Successfully parsed 1528 lines"),
         (logging.INFO, "Found 102 training entries"),
         (logging.INFO, "Found 34 validation entries"),
     ]

tracking/pyproject.toml

@@ -9,7 +9,7 @@ description = "A training log parser for the Mozilla translation ML tool"
 requires-python = ">=3.10"
 dependencies = [
   "PyYAML>=6.0.1",
-  "wandb==0.16.1",
+  "wandb==0.17.2",
   "taskcluster>=56.0.3,<57.0.0",
   # numpy 2.0 fails with wandb
   "numpy<2.0.0"

tracking/translations_parser/data.py

@@ -143,6 +143,5 @@ class TrainingLog:
     configuration: dict
     training: List[TrainingEpoch]
     validation: List[ValidationEpoch]
-    # Dict of log lines indexed by their header (e.g. marian, data, memory)
-    logs: dict
+    logs: List[str]
     run_date: datetime | None

tracking/translations_parser/parser.py

@@ -5,7 +5,7 @@
 from collections.abc import Iterable, Iterator, Sequence
 from datetime import datetime
 from itertools import tee
-from typing import Callable, DefaultDict
+from typing import Callable, DefaultDict, List
 
 import yaml
 
@@ -53,7 +53,7 @@ def __init__(
         self._current_index = 0
         self.parsed = False
         self.config: dict = {}
-        self.indexed_logs: DefaultDict[str, list] = defaultdict(list)
+        self.parsed_logs: List[str] = []
         # Optional list of Metric published earlier to the parsing
         self.metrics = metrics
         self.training: list[TrainingEpoch] = []
@@ -195,19 +195,23 @@ def _iter_log_entries(self) -> Iterator[tuple[list[tuple[str]], str]]:
                 self.run_date = self.get_timestamp(headers)
             text = line[position:]
 
-            def _join(iterable):
-                if not iterable:
-                    return "_"
-                if isinstance(iterable[0], str):
-                    return "_".join(iterable)
-                return _join([_join(item) for item in iterable])
+            def _join(seq):
+                if not seq:
+                    return None
+                if isinstance(seq[0], str):
+                    return "_".join(seq)
+                return _join([_join(item) for item in seq])
 
             # Record logs depending on Marian headers
+            tag = None
             if len(headers) >= 2:
-                # First is task timestamp, second is marian timestamp
+                # The 2 first headers are ignored (task timestamp, then marian timestamp)
                 _, _, *marian_tags = headers
                 tag = _join(marian_tags)
-                self.indexed_logs[tag].append(text)
+            if tag:
+                self.parsed_logs.append(f"[tag] {text}")
+            else:
+                self.parsed_logs.append(text)
 
             yield headers, text
 
@@ -316,13 +320,6 @@ def parse(self) -> None:
         for publisher in self.publishers:
             publisher.close()
 
-    @property
-    def logs_str(self) -> str:
-        return "\n".join(
-            "".join(f"[{key}] {val}\n" for val in values)
-            for key, values in self.indexed_logs.items()
-        )
-
     @property
     def output(self) -> TrainingLog:
         if not self.parsed:
@@ -332,7 +329,7 @@ def output(self) -> TrainingLog:
             configuration=self.config,
             training=self.training,
             validation=list(self.validation),
-            logs=self.indexed_logs,
+            logs=self.parsed_logs,
         )
 
     def run(self) -> None:
@@ -343,7 +340,7 @@ def run(self) -> None:
             # A StopIteration can be raised if some required lines are never found.
             raise ValueError("Logs file ended up unexpectedly")
 
-        count = sum(len(vals) for vals in self.indexed_logs.values())
+        count = len(self.parsed_logs)
         logger.info(f"Successfully parsed {count} lines")
         logger.info(f"Found {len(self.training)} training entries")
         logger.info(f"Found {len(self.validation)} validation entries")