Fix COMET_EXPERIMENT_KEY environment variable usage in comet logger (#4230)

* Fix  COMET_EXPERIMENT_KEY environment variable usage

* Remove unused arg

* Update comet.py

* Add test by Lothiraldan

* remove blank

Co-authored-by: chaton <[email protected]>
Co-authored-by: Nicki Skafte <[email protected]>
Co-authored-by: Sean Naren <[email protected]>

Author: 4 people

pytorch_lightning/loggers/comet.py

@@ -187,7 +187,6 @@ def experiment(self):
 
         if self._future_experiment_key is not None:
             os.environ["COMET_EXPERIMENT_KEY"] = self._future_experiment_key
-            self._future_experiment_key = None
 
         try:
             if self.mode == "online":
@@ -212,7 +211,9 @@ def experiment(self):
                     **self._kwargs,
                 )
         finally:
-            os.environ.pop("COMET_EXPERIMENT_KEY", None)
+            if self._future_experiment_key is not None:
+                os.environ.pop("COMET_EXPERIMENT_KEY")
+                self._future_experiment_key = None
 
         if self._experiment_name:
             self._experiment.set_name(self._experiment_name)
@@ -278,6 +279,9 @@ def version(self) -> str:
         if self._experiment_key is not None:
             return self._experiment_key
 
+        if "COMET_EXPERIMENT_KEY" in os.environ:
+            return os.environ["COMET_EXPERIMENT_KEY"]
+
         if self._future_experiment_key is not None:
             return self._future_experiment_key
 

tests/loggers/test_comet.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
-from unittest.mock import patch
+from unittest.mock import patch, DEFAULT
 
 import pytest
 
@@ -99,6 +99,37 @@ def test_comet_logger_experiment_name(comet):
         comet_experiment().set_name.assert_called_once_with(experiment_name)
 
 
+@patch('pytorch_lightning.loggers.comet.comet_ml')
+def test_comet_logger_manual_experiment_key(comet):
+    """Test that Comet Logger respects manually set COMET_EXPERIMENT_KEY."""
+
+    api_key = "key"
+    experiment_key = "96346da91469407a85641afe5766b554"
+
+    instantation_environ = {}
+
+    def save_os_environ(*args, **kwargs):
+        nonlocal instantation_environ
+        instantation_environ = os.environ.copy()
+
+        return DEFAULT
+
+    # Test api_key given
+    with patch.dict(os.environ, {"COMET_EXPERIMENT_KEY": experiment_key}):
+        with patch('pytorch_lightning.loggers.comet.CometExperiment', side_effect=save_os_environ) as comet_experiment:
+            logger = CometLogger(api_key=api_key)
+
+            assert logger.version == experiment_key
+
+            assert logger._experiment is None
+
+            _ = logger.experiment
+
+            comet_experiment.assert_called_once_with(api_key=api_key, project_name=None)
+
+    assert instantation_environ["COMET_EXPERIMENT_KEY"] == experiment_key
+
+
 @patch('pytorch_lightning.loggers.comet.CometOfflineExperiment')
 @patch('pytorch_lightning.loggers.comet.comet_ml')
 def test_comet_logger_dirs_creation(comet, comet_experiment, tmpdir, monkeypatch):