autogen/test/test_training_log.py
Chi Wang 776aa55189
V0.2.2 (#19)
* v0.2.2

separate the HPO part into the module flaml.tune
enhanced implementation of FLOW^2, CFO and BlendSearch
support parallel tuning using ray tune
add support for sample_weight and generic fit arguments
enable mlflow logging

Co-authored-by: Chi Wang (MSR) <chiw@microsoft.com>
Co-authored-by: qingyun-wu <qw2ky@virginia.edu>
2021-02-05 21:41:14 -08:00

42 lines
1.3 KiB
Python

import os
import unittest
from tempfile import TemporaryDirectory
from sklearn.datasets import load_boston
from flaml import AutoML
from flaml.training_log import training_log_reader
class TestTrainingLog(unittest.TestCase):
def test_training_log(self):
with TemporaryDirectory() as d:
filename = os.path.join(d, 'test_training_log.log')
# Run a simple job.
automl_experiment = AutoML()
automl_settings = {
"time_budget": 2,
"metric": 'mse',
"task": 'regression',
"log_file_name": filename,
"log_training_metric": True,
"mem_thres": 1024*1024,
"n_jobs": 1,
"model_history": True
}
X_train, y_train = load_boston(return_X_y=True)
automl_experiment.fit(X_train=X_train, y_train=y_train,
**automl_settings)
# Check if the training log file is populated.
self.assertTrue(os.path.exists(filename))
with training_log_reader(filename) as reader:
count = 0
for record in reader.records():
print(record)
count += 1
self.assertGreater(count, 0)