From 54e78bde8552b9e83f3976a26fac7981a84ac6ca Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9E=97SO?= <142557582+Linxiushen@users.noreply.github.com> Date: Mon, 28 Sep 2026 01:45:01 +0800 Subject: [PATCH 1/2] fix(automl): preserve record IDs when appending training logs --- flaml/automl/training_log.py | 6 ++ test/automl/test_training_log_append.py | 82 +++++++++++++++++++++++++ 2 files changed, 88 insertions(+) create mode 100644 test/automl/test_training_log_append.py diff --git a/flaml/automl/training_log.py b/flaml/automl/training_log.py index 578b8de298..3d4eef7577 100644 --- a/flaml/automl/training_log.py +++ b/flaml/automl/training_log.py @@ -65,6 +65,12 @@ def open(self): self.file = open(self.output_filename, "w") def append_open(self): + try: + with training_log_reader(self.output_filename) as reader: + for record in reader.records(): + self.current_record_id = max(self.current_record_id, record.record_id + 1) + except FileNotFoundError: + pass self.file = open(self.output_filename, "a") def append( diff --git a/test/automl/test_training_log_append.py b/test/automl/test_training_log_append.py new file mode 100644 index 0000000000..71777f9e78 --- /dev/null +++ b/test/automl/test_training_log_append.py @@ -0,0 +1,82 @@ +import json + +import pytest + +from flaml import AutoML +from flaml.automl.training_log import training_log_reader, training_log_writer + + +def append_trial(writer, n_estimators): + writer.append(1, None, 0.1, 0.1, 1.0 / n_estimators, {"n_estimators": n_estimators}, "rf", 100) + + +def test_appended_trial_can_be_loaded_from_checkpoint(tmp_path): + filename = tmp_path / "training.log" + for count, append in ((4, False), (8, True), (16, True)): + with training_log_writer(filename, append=append) as writer: + append_trial(writer, count) + writer.checkpoint() + + records = [json.loads(line) for line in filename.read_text().splitlines()] + checkpoint_id = records[-1]["curr_best_record_id"] + estimator = AutoML().get_estimator_from_log(str(filename), checkpoint_id, "classification") + assert estimator.n_estimators == 16 + assert [record["record_id"] for record in records if "record_id" in record] == [0, 1, 2] + assert [record["curr_best_record_id"] for record in records if "curr_best_record_id" in record] == [0, 1, 2] + + +@pytest.mark.parametrize("existing_ids", [[], [0], [2, 5], [5, 2], [0, 1, 0]]) +def test_append_uses_next_unused_record_id(tmp_path, existing_ids): + filename = tmp_path / "training.log" + with training_log_writer(filename) as writer: + for record_id in existing_ids: + writer.current_record_id = record_id + append_trial(writer, 4) + writer.checkpoint() + original = filename.read_bytes() + + with training_log_writer(filename, append=True) as writer: + append_trial(writer, 8) + append_trial(writer, 16) + writer.checkpoint() + + assert filename.read_bytes().startswith(original) + expected = max(existing_ids, default=-1) + 1 + with training_log_reader(filename) as reader: + records = list(reader.records()) + assert [record.record_id for record in records] == existing_ids + [expected, expected + 1] + with training_log_reader(filename) as reader: + assert reader.get_record(expected).config == {"n_estimators": 8} + + +def test_append_creates_missing_log(tmp_path): + filename = tmp_path / "new.log" + with training_log_writer(filename, append=True) as writer: + append_trial(writer, 4) + + with training_log_reader(filename) as reader: + assert reader.get_record(0).config == {"n_estimators": 4} + + +def test_appended_checkpoint_keeps_current_run_best(tmp_path): + filename = tmp_path / "training.log" + with training_log_writer(filename) as writer: + append_trial(writer, 16) + writer.checkpoint() + with training_log_writer(filename, append=True) as writer: + append_trial(writer, 4) + writer.checkpoint() + + records = [json.loads(line) for line in filename.read_text().splitlines()] + assert records[-1] == {"curr_best_record_id": 1} + with training_log_reader(filename) as reader: + assert reader.get_record(1).config == {"n_estimators": 4} + + +def test_append_skips_checkpoint_only_log(tmp_path): + filename = tmp_path / "training.log" + filename.write_text('{"curr_best_record_id": 10}\n') + with training_log_writer(filename, append=True) as writer: + append_trial(writer, 4) + with training_log_reader(filename) as reader: + assert reader.get_record(0).config == {"n_estimators": 4} From a832a8a40d12d3b77d4203b246f28be4a4941c57 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9E=97SO?= <142557582+Linxiushen@users.noreply.github.com> Date: Mon, 28 Sep 2026 05:50:15 +0800 Subject: [PATCH 2/2] fix(automl): serialize training log writers and terminate tails --- flaml/automl/automl.py | 12 ++- flaml/automl/training_log.py | 48 +++++++-- pyproject.toml | 2 + test/automl/test_training_log_append.py | 128 ++++++++++++++++++++++++ 4 files changed, 176 insertions(+), 14 deletions(-) diff --git a/flaml/automl/automl.py b/flaml/automl/automl.py index 1a7bee036a..3ce2b88315 100644 --- a/flaml/automl/automl.py +++ b/flaml/automl/automl.py @@ -286,8 +286,10 @@ def custom_metric( mlflow_exp_name: str, default=None | The name of the mlflow experiment. This should be specified if enable mlflow autologging on Spark. Otherwise it will log all the results into the experiment of the same name as the basename of main entry file. - append_log: boolean, default=False | Whetehr to directly append the log - records to the input log file if it exists. + append_log: boolean, default=False | Whether to directly append the log + records to the input log file if it exists. Runs sharing a log file + wait for its writer to close; this wait counts toward time_budget. + Use different log files for concurrent runs. auto_augment: boolean, default=True | Whether to automatically augment rare classes. resampler: object, default=None | An imbalanced-learn-compatible resampler @@ -2114,8 +2116,10 @@ def custom_metric( mlflow_exp_name: str, default=None | The name of the mlflow experiment. This should be specified if enable mlflow autologging on Spark. Otherwise it will log all the results into the experiment of the same name as the basename of main entry file. - append_log: boolean, default=False | Whetehr to directly append the log - records to the input log file if it exists. + append_log: boolean, default=False | Whether to directly append the log + records to the input log file if it exists. Runs sharing a log file + wait for its writer to close; this wait counts toward time_budget. + Use different log files for concurrent runs. auto_augment: boolean, default=True | Whether to automatically augment rare classes. min_sample_size: int, default=MIN_SAMPLE_TRAIN | the minimal sample diff --git a/flaml/automl/training_log.py b/flaml/automl/training_log.py index 3d4eef7577..45e97a3fdc 100644 --- a/flaml/automl/training_log.py +++ b/flaml/automl/training_log.py @@ -5,6 +5,7 @@ import json import logging +import os from contextlib import contextmanager from typing import IO @@ -56,22 +57,43 @@ class TrainingLogWriter: def __init__(self, output_filename: str): self.output_filename = output_filename self.file = None + self._lock = None self.current_best_loss_record_id = None self.current_best_loss = float("+inf") self.current_sample_size = None self.current_record_id = 0 def open(self): - self.file = open(self.output_filename, "w") + self._open(append=False) def append_open(self): + self._open(append=True) + + def _open(self, append): + from filelock import FileLock + + self.close() + self._lock = FileLock(f"{os.path.realpath(self.output_filename)}.lock") try: - with training_log_reader(self.output_filename) as reader: - for record in reader.records(): - self.current_record_id = max(self.current_record_id, record.record_id + 1) - except FileNotFoundError: - pass - self.file = open(self.output_filename, "a") + self._lock.acquire() + if append: + try: + with training_log_reader(self.output_filename) as reader: + for record in reader.records(): + self.current_record_id = max(self.current_record_id, record.record_id + 1) + except FileNotFoundError: + pass + self.file = open(self.output_filename, "a" if append else "w") + if append: + with open(self.output_filename, "rb") as existing: + existing.seek(0, os.SEEK_END) + if existing.tell(): + existing.seek(-1, os.SEEK_END) + if existing.read(1) not in (b"\n", b"\r"): + self.file.write("\n") + except BaseException: + self.close() + raise def append( self, @@ -125,9 +147,14 @@ def checkpoint(self): self.file.flush() def close(self): - if self.file is not None: - self.file.close() - self.file = None # for pickle + try: + if self.file is not None: + self.file.close() + finally: + self.file = None # for pickle + lock, self._lock = self._lock, None + if lock is not None: + lock.release() class TrainingLogReader: @@ -164,6 +191,7 @@ def get_record(self, record_id) -> TrainingLogRecord: @contextmanager def training_log_writer(filename: str, append: bool = False): + """Open a training log, holding its sidecar file lock until the writer closes.""" try: w = TrainingLogWriter(filename) if not append: diff --git a/pyproject.toml b/pyproject.toml index 537282d6d5..7a6e2be194 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,6 +30,7 @@ Homepage = "https://github.com/microsoft/FLAML" [project.optional-dependencies] automl = [ + "filelock>=3.0.0", "lightgbm>=2.3.1", "xgboost>=0.90,<3.0.0", "scipy>=1.4.1", @@ -46,6 +47,7 @@ spark = [ "joblib<=1.3.2", ] test = [ + "filelock>=3.0.0", "numpy>=1.17,<2.0.0; python_version<'3.13' and (sys_platform!='win32' or platform_machine!='ARM64')", "numpy>=1.17; python_version>='3.13' or (sys_platform=='win32' and platform_machine=='ARM64')", "jupyter", diff --git a/test/automl/test_training_log_append.py b/test/automl/test_training_log_append.py index 71777f9e78..36b1fce747 100644 --- a/test/automl/test_training_log_append.py +++ b/test/automl/test_training_log_append.py @@ -1,4 +1,6 @@ import json +import multiprocessing +import pickle import pytest @@ -10,6 +12,16 @@ def append_trial(writer, n_estimators): writer.append(1, None, 0.1, 0.1, 1.0 / n_estimators, {"n_estimators": n_estimators}, "rf", 100) +def append_in_process(filename, n_estimators, started, opened, release=None, append=True): + started.set() + with training_log_writer(filename, append=append) as writer: + opened.set() + if release is not None and not release.wait(60): + raise TimeoutError("Parent did not release the training log writer") + append_trial(writer, n_estimators) + writer.checkpoint() + + def test_appended_trial_can_be_loaded_from_checkpoint(tmp_path): filename = tmp_path / "training.log" for count, append in ((4, False), (8, True), (16, True)): @@ -80,3 +92,119 @@ def test_append_skips_checkpoint_only_log(tmp_path): append_trial(writer, 4) with training_log_reader(filename) as reader: assert reader.get_record(0).config == {"n_estimators": 4} + + +@pytest.mark.parametrize("checkpoint_tail", [False, True]) +@pytest.mark.parametrize("line_ending", [b"", b"\n", b"\r\n", b"\r"]) +def test_append_preserves_valid_record_and_checkpoint_tails(tmp_path, checkpoint_tail, line_ending): + filename = tmp_path / "training.log" + with training_log_writer(filename) as writer: + append_trial(writer, 4) + if checkpoint_tail: + writer.checkpoint() + original = filename.read_bytes().rstrip(b"\r\n") + line_ending + filename.write_bytes(original) + + with training_log_writer(filename, append=True) as writer: + append_trial(writer, 8) + writer.checkpoint() + + assert filename.read_bytes().startswith(original) + with training_log_reader(filename) as reader: + assert [record.record_id for record in reader.records()] == [0, 1] + checkpoint_id = json.loads(filename.read_text().splitlines()[-1])["curr_best_record_id"] + estimator = AutoML().get_estimator_from_log(str(filename), checkpoint_id, "classification") + assert estimator.n_estimators == 8 + + +@pytest.mark.parametrize( + "existing_log, same_log, first_append", + [(False, True, True), (True, True, True), (True, True, False), (True, False, True)], +) +def test_process_writers_serialize_only_for_the_same_log(tmp_path, existing_log, same_log, first_append): + filename = tmp_path / "training.log" + if existing_log: + with training_log_writer(filename) as writer: + append_trial(writer, 4) + second_filename = filename if same_log else tmp_path / "other.log" + context = multiprocessing.get_context("spawn") + first_started, first_opened, release_first = (context.Event() for _ in range(3)) + second_started, second_opened = (context.Event() for _ in range(2)) + first = context.Process( + target=append_in_process, args=(filename, 8, first_started, first_opened, release_first, first_append) + ) + second = context.Process(target=append_in_process, args=(second_filename, 16, second_started, second_opened)) + processes = [] + try: + first.start() + processes.append(first) + assert first_opened.wait(60), "First writer did not open the log" + second.start() + processes.append(second) + assert second_started.wait(60), "Second writer did not start" + opened_before_release = second_opened.wait(3 if same_log else 60) + release_first.set() + for process in processes: + process.join(timeout=60) + assert process.exitcode == 0 + finally: + release_first.set() + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=10) + + with training_log_reader(filename) as reader: + records = list(reader.records()) + expected_counts = ([4] if existing_log and first_append else []) + [8] + ([16] if same_log else []) + assert [record.record_id for record in records] == list(range(len(expected_counts))) + assert [record.config["n_estimators"] for record in records] == expected_counts + assert opened_before_release is not same_log + if not same_log: + with training_log_reader(second_filename) as reader: + assert reader.get_record(0).config == {"n_estimators": 16} + + +@pytest.mark.parametrize("append", [False, True]) +def test_writer_releases_lock_after_context_error(tmp_path, append): + from filelock import FileLock + + filename = tmp_path / "training.log" + with pytest.raises(RuntimeError, match="training failed"): + with training_log_writer(filename, append=append) as writer: + append_trial(writer, 4) + raise RuntimeError("training failed") + + writer.close() + assert pickle.loads(pickle.dumps(writer)).file is None + with FileLock(f"{filename}.lock", timeout=0): + with training_log_reader(filename) as reader: + assert reader.get_record(0).config == {"n_estimators": 4} + + +def test_append_releases_lock_after_invalid_log(tmp_path): + from filelock import FileLock + + filename = tmp_path / "training.log" + original = '{"record_id":' + filename.write_text(original) + with pytest.raises(json.JSONDecodeError): + with training_log_writer(filename, append=True): + pytest.fail("An invalid log should not be opened for append") + + assert filename.read_text() == original + with FileLock(f"{filename}.lock", timeout=0): + pass + + +def test_writer_releases_lock_after_open_error(tmp_path): + from filelock import FileLock + + filename = tmp_path / "training.log" + filename.mkdir() + with pytest.raises(OSError): + with training_log_writer(filename): + pytest.fail("A directory should not be opened as a log file") + + with FileLock(f"{filename}.lock", timeout=0): + pass