From 0489e31b02e3397e460f2347f445a4b155c111f9 Mon Sep 17 00:00:00 2001
From: Nic Eggert <nic@eggert.io>
Date: Fri, 6 Dec 2019 23:24:59 -0600
Subject: [PATCH] Fix CometML tests (#585)

* monkeypatch atexit.register to fix problem with cometml logging

* Use experiment id for version in cometml
---
 pytorch_lightning/logging/comet.py |  7 +------
 tests/test_logging.py              | 16 ++++++++++++++--
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/pytorch_lightning/logging/comet.py b/pytorch_lightning/logging/comet.py
index 3bfcf9b1..3fef13a2 100644
--- a/pytorch_lightning/logging/comet.py
+++ b/pytorch_lightning/logging/comet.py
@@ -167,9 +167,4 @@ class CometLogger(LightningLoggerBase):
 
     @property
     def version(self):
-        if self.project_name and self.rest_api_key:
-            # Determines the number of experiments in this project, and returns the next integer as the version number
-            num_exps = len(self.comet_api.get_experiments(self.workspace, self.project_name))
-            return num_exps + 1
-        else:
-            return None
+        return self.experiment.id
diff --git a/tests/test_logging.py b/tests/test_logging.py
index 37618a56..8f90fe41 100644
--- a/tests/test_logging.py
+++ b/tests/test_logging.py
@@ -106,8 +106,14 @@ def test_mlflow_pickle(tmpdir):
     trainer2.logger.log_metrics({"acc": 1.0})
 
 
-def test_comet_logger(tmpdir):
+def test_comet_logger(tmpdir, monkeypatch):
     """Verify that basic functionality of Comet.ml logger works."""
+
+    # prevent comet logger from trying to print at exit, since
+    # pytest's stdout/stderr redirection breaks it
+    import atexit
+    monkeypatch.setattr(atexit, "register", lambda _: None)
+
     tutils.reset_seed()
 
     try:
@@ -140,8 +146,14 @@ def test_comet_logger(tmpdir):
     assert result == 1, "Training failed"
 
 
-def test_comet_pickle(tmpdir):
+def test_comet_pickle(tmpdir, monkeypatch):
     """Verify that pickling trainer with comet logger works."""
+
+    # prevent comet logger from trying to print at exit, since
+    # pytest's stdout/stderr redirection breaks it
+    import atexit
+    monkeypatch.setattr(atexit, "register", lambda _: None)
+
     tutils.reset_seed()
 
     try: