From 7126313b29bfa9002410c3f4e575d163b9aed05a Mon Sep 17 00:00:00 2001 From: Will Sackfield Date: Mon, 25 Nov 2024 17:32:25 -0500 Subject: [PATCH 1/2] Add test for cog stream draining file logging MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * This tests exposes the issue that cog stream drainer can’t drain due to file logging in the logging module. --- .../cog.yaml | 5 +++++ .../predict.py | 22 +++++++++++++++++++ .../test_integration/test_predict.py | 16 ++++++++++++++ 3 files changed, 43 insertions(+) create mode 100644 test-integration/test_integration/fixtures/file-logging-stream-drain-project/cog.yaml create mode 100644 test-integration/test_integration/fixtures/file-logging-stream-drain-project/predict.py 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, + ) From 814e72a96af73497e0535b73e50bdd6752ea145d Mon Sep 17 00:00:00 2001 From: Will Sackfield Date: Tue, 26 Nov 2024 17:09:54 -0500 Subject: [PATCH 2/2] Add test for handling tqdm stream redirection --- .../cog.yaml | 2 +- .../predict.py | 26 ++++++------------- .../test_integration/test_predict.py | 1 - 3 files changed, 9 insertions(+), 20 deletions(-) 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 index 31d29ef04b..4d02bd8d14 100644 --- 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 @@ -1,5 +1,5 @@ build: python_version: "3.11" python_packages: - - "pillow==8.3.2" + - "tqdm==4.66.6" 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 index 5ce00f182c..4eb76c5eba 100644 --- 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 @@ -1,22 +1,12 @@ -from cog import BasePredictor, File -import logging -import os -from PIL import Image +from cog import BasePredictor +import tqdm +import time 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 [] + def predict(self) -> str: + print("Starting Predict:") + for i in tqdm.tqdm(range(10), total=10): + time.sleep(i) + return "Hello World" diff --git a/test-integration/test_integration/test_predict.py b/test-integration/test_integration/test_predict.py index e417d56537..a25f4e2aa6 100644 --- a/test-integration/test_integration/test_predict.py +++ b/test-integration/test_integration/test_predict.py @@ -334,7 +334,6 @@ def test_file_logging_stream_drain(tmpdir_factory): cmd, cwd=out_dir, check=True, - capture_output=True, text=True, timeout=DEFAULT_TIMEOUT, )