diff --git a/test-integration/test_integration/fixtures/file-logging-stream-drain-project/cog.yaml b/test-integration/test_integration/fixtures/file-logging-stream-drain-project/cog.yaml new file mode 100644 index 0000000000..31d29ef04b --- /dev/null +++ b/test-integration/test_integration/fixtures/file-logging-stream-drain-project/cog.yaml @@ -0,0 +1,5 @@ +build: + python_version: "3.11" + python_packages: + - "pillow==8.3.2" +predict: "predict.py:Predictor" diff --git a/test-integration/test_integration/fixtures/file-logging-stream-drain-project/predict.py b/test-integration/test_integration/fixtures/file-logging-stream-drain-project/predict.py new file mode 100644 index 0000000000..5ce00f182c --- /dev/null +++ b/test-integration/test_integration/fixtures/file-logging-stream-drain-project/predict.py @@ -0,0 +1,22 @@ +from cog import BasePredictor, File +import logging +import os +from PIL import Image + + +class Predictor(BasePredictor): + def setup(self): + logFormatter = logging.Formatter("%(asctime)s [%(threadName)-12.12s] [%(levelname)-5.5s] %(message)s") + rootLogger = logging.getLogger() + + fileHandler = logging.FileHandler("{0}/{1}.log".format(os.path.dirname(__file__), "mylog.log")) + fileHandler.setFormatter(logFormatter) + rootLogger.addHandler(fileHandler) + + consoleHandler = logging.StreamHandler() + consoleHandler.setFormatter(logFormatter) + rootLogger.addHandler(consoleHandler) + + def predict(self) -> list[Image]: + logging.info("Do some logging.") + return [] diff --git a/test-integration/test_integration/test_predict.py b/test-integration/test_integration/test_predict.py index 4065b69927..e417d56537 100644 --- a/test-integration/test_integration/test_predict.py +++ b/test-integration/test_integration/test_predict.py @@ -322,3 +322,19 @@ def test_predict_with_subprocess_in_setup(fixture_name): assert response.status_code == 200, str(response) assert busy_count < 10 + + +def test_file_logging_stream_drain(tmpdir_factory): + project_dir = Path(__file__).parent / "fixtures/file-logging-stream-drain-project" + out_dir = pathlib.Path(tmpdir_factory.mktemp("project")) + shutil.copytree(project_dir, out_dir, dirs_exist_ok=True) + cmd = ["cog", "--debug", "predict"] + + subprocess.run( + cmd, + cwd=out_dir, + check=True, + capture_output=True, + text=True, + timeout=DEFAULT_TIMEOUT, + )