Files
wehub-resource-sync 593b94c120
pytest / Unit Tests (push) Has been cancelled
pytest / Integration (integration_tests_a) (push) Has been cancelled
pytest / Integration (integration_tests_b) (push) Has been cancelled
pytest / Integration (integration_tests_c) (push) Has been cancelled
pytest / Integration (integration_tests_d) (push) Has been cancelled
pytest / Integration (integration_tests_e) (push) Has been cancelled
pytest / Integration (integration_tests_f) (push) Has been cancelled
pytest / Integration (integration_tests_g) (push) Has been cancelled
pytest / Integration (integration_tests_h) (push) Has been cancelled
pytest / Integration (integration_tests_i) (push) Has been cancelled
pytest / Integration (integration_tests_j) (push) Has been cancelled
pytest / Distributed (distributed_a) (push) Has been cancelled
pytest / Distributed (distributed_b) (push) Has been cancelled
pytest / Distributed (distributed_c) (push) Has been cancelled
pytest / Distributed (distributed_d) (push) Has been cancelled
pytest / Distributed (distributed_e) (push) Has been cancelled
pytest / Distributed (distributed_f) (push) Has been cancelled
pytest / Minimal Install (push) Has been cancelled
pytest / Event File (push) Has been cancelled
pytest (slow) / py-slow (push) Has been cancelled
Publish JSON Schema / publish-schema (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:49:20 +08:00

80 lines
2.7 KiB
Python

# Tests the following end-to-end:
#
# 1. Comet is imported
# 2. Conflicting modules (i.e., TensorFlow) are not imported
# 3. Overridden methods are called (train_init, train_model, etc.) and run without error
#
# This test runs in an isolated environment to ensure TensorFlow imports are not leaked
# from previous tests.
import argparse
import os
import sys
import tempfile
from unittest.mock import Mock, patch
# Comet must be imported before the libraries it wraps
import comet_ml # noqa
from ludwig.api import LudwigModel
from ludwig.constants import BATCH_SIZE, TRAINER
from ludwig.contribs.comet import CometCallback
# Bad key will ensure Comet is initialized, but nothing is uploaded externally.
os.environ["COMET_API_KEY"] = "key"
# Add tests dir to the import path
PATH_HERE = os.path.abspath(os.path.dirname(__file__))
PATH_ROOT = os.path.join(PATH_HERE, "..", "..", "..")
sys.path.insert(0, os.path.abspath(PATH_ROOT))
from tests.integration_tests.utils import category_feature, generate_data, image_feature # noqa
parser = argparse.ArgumentParser()
parser.add_argument("--csv-filename", required=True)
def run(csv_filename):
with tempfile.TemporaryDirectory() as tmpdir:
# Image Inputs
image_dest_folder = os.path.join(tmpdir, "generated_images")
# Inputs & Outputs
input_features = [image_feature(folder=image_dest_folder)]
output_features = [category_feature(output_feature=True)]
data_csv = generate_data(input_features, output_features, csv_filename)
config = {
"input_features": input_features,
"output_features": output_features,
"combiner": {"type": "concat", "output_size": 14},
TRAINER: {"epochs": 2, BATCH_SIZE: 128},
}
callback = CometCallback()
model = LudwigModel(config, callbacks=[callback])
# Wrap these methods so we can check that they were called
callback.on_train_init = Mock(side_effect=callback.on_train_init)
callback.on_train_start = Mock(side_effect=callback.on_train_start)
with patch("comet_ml.Experiment.log_asset_data") as mock_log_asset_data:
# Training with csv
_, _, _ = model.train(dataset=data_csv, output_directory=os.path.join(tmpdir, "output"))
model.predict(dataset=data_csv)
# Verify that the experiment was created successfully
assert callback.cometml_experiment is not None
# Check that these methods were called at least once
callback.on_train_init.assert_called()
callback.on_train_start.assert_called()
# Check that we ran `train_model`, which calls into `log_assert_data`, successfully
mock_log_asset_data.assert_called()
if __name__ == "__main__":
args = parser.parse_args()
run(args.csv_filename)