feature/comet-logger

src/lightning/pytorch/loggers/comet.py

@@ -19,14 +19,14 @@
 import logging
 import os
 from argparse import Namespace
-from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Union
+from typing import Any, Dict, Literal, Mapping, Optional, TYPE_CHECKING, Union
 
 from lightning_utilities.core.imports import RequirementCache
 from torch import Tensor
 from torch.nn import Module
 from typing_extensions import override
 
-from lightning.fabric.utilities.logger import _add_prefix, _convert_params, _flatten_dict
+from lightning.fabric.utilities.logger import _add_prefix, _convert_params
 from lightning.pytorch.loggers.logger import Logger, rank_zero_experiment
 from lightning.pytorch.utilities.exceptions import MisconfigurationException
 from lightning.pytorch.utilities.rank_zero import rank_zero_only
@@ -35,7 +35,7 @@
     from comet_ml import ExistingExperiment, Experiment, OfflineExperiment
 
 log = logging.getLogger(__name__)
-_COMET_AVAILABLE = RequirementCache("comet-ml>=3.31.0", module="comet_ml")
+_COMET_AVAILABLE = RequirementCache("comet-ml>=3.44.4", module="comet_ml")
 
 
 class CometLogger(Logger):
@@ -64,7 +64,6 @@ class CometLogger(Logger):
             workspace=os.environ.get("COMET_WORKSPACE"),  # Optional
             save_dir=".",  # Optional
             project_name="default_project",  # Optional
-            rest_api_key=os.environ.get("COMET_REST_API_KEY"),  # Optional
             experiment_key=os.environ.get("COMET_EXPERIMENT_KEY"),  # Optional
             experiment_name="lightning_logs",  # Optional
         )
@@ -81,7 +80,6 @@ class CometLogger(Logger):
             save_dir=".",
             workspace=os.environ.get("COMET_WORKSPACE"),  # Optional
             project_name="default_project",  # Optional
-            rest_api_key=os.environ.get("COMET_REST_API_KEY"),  # Optional
             experiment_name="lightning_logs",  # Optional
         )
         trainer = Trainer(logger=comet_logger)
@@ -166,24 +164,22 @@ def __init__(self, *args, **kwarg):
         - `Comet Documentation <https://www.comet.com/docs/v2/integrations/ml-frameworks/pytorch-lightning/>`__
 
     Args:
-        api_key: Required in online mode. API key, found on Comet.ml. If not given, this
+        api_key: Required in online mode. API key, found on Comet.com. If not given, this
             will be loaded from the environment variable COMET_API_KEY or ~/.comet.config
             if either exists.
         save_dir: Required in offline mode. The path for the directory to save local
             comet logs. If given, this also sets the directory for saving checkpoints.
         project_name: Optional. Send your experiment to a specific project.
-            Otherwise will be sent to Uncategorized Experiments.
-            If the project name does not already exist, Comet.ml will create a new project.
-        rest_api_key: Optional. Rest API key found in Comet.ml settings.
-            This is used to determine version number
-        experiment_name: Optional. String representing the name for this particular experiment on Comet.ml.
+            Otherwise, will be sent to Uncategorized Experiments.
+            If the project name does not already exist, Comet.com will create a new project.
+        experiment_name: Optional. String representing the name for this particular experiment on Comet.com.
         experiment_key: Optional. If set, restores from existing experiment.
         offline: If api_key and save_dir are both given, this determines whether
             the experiment will be in online or offline mode. This is useful if you use
             save_dir to control the checkpoints directory and have a ~/.comet.config
             file but still want to run offline experiments.
         prefix: A string to put at the beginning of metric keys.
-        \**kwargs: Additional arguments like `workspace`, `log_code`, etc. used by
+        **kwargs: Additional arguments like `workspace`, `log_code`, etc. used by
             :class:`CometExperiment` can be passed as keyword arguments in this logger.
 
     Raises:
@@ -201,7 +197,6 @@ def __init__(
         api_key: Optional[str] = None,
         save_dir: Optional[str] = None,
         project_name: Optional[str] = None,
-        rest_api_key: Optional[str] = None,
         experiment_name: Optional[str] = None,
         experiment_key: Optional[str] = None,
         offline: bool = False,
@@ -210,16 +205,22 @@ def __init__(
     ):
         if not _COMET_AVAILABLE:
             raise ModuleNotFoundError(str(_COMET_AVAILABLE))
-        super().__init__()
-        self._experiment = None
+
         self._save_dir: Optional[str]
-        self.rest_api_key: Optional[str]
+        self.api_key: str
+        self.mode: Literal["online", "offline"]
+
+        super().__init__()
 
         # needs to be set before the first `comet_ml` import
+        # because comet_ml imported after another machine learning libraries (Torch)
         os.environ["COMET_DISABLE_AUTO_LOGGING"] = "1"
 
         import comet_ml
 
+        comet_experiment = Union[comet_ml.Experiment, comet_ml.ExistingExperiment, comet_ml.OfflineExperiment]
+        self._experiment: Optional[comet_experiment] = None
+
         # Determine online or offline mode based on which arguments were passed to CometLogger
         api_key = api_key or comet_ml.config.get_api_key(None, comet_ml.config.get_config())
 
@@ -241,21 +242,10 @@ def __init__(
         log.info(f"CometLogger will be initialized in {self.mode} mode")
 
         self._project_name: Optional[str] = project_name
-        self._experiment_key: Optional[str] = experiment_key
+        self._experiment_key: str = experiment_key or os.environ.get("COMET_EXPERIMENT_KEY") or comet_ml.generate_guid()
         self._experiment_name: Optional[str] = experiment_name
         self._prefix: str = prefix
-        self._kwargs: Any = kwargs
-        self._future_experiment_key: Optional[str] = None
-
-        if rest_api_key is not None:
-            from comet_ml.api import API
-
-            # Comet.ml rest API, used to determine version number
-            self.rest_api_key = rest_api_key
-            self.comet_api = API(self.rest_api_key)
-        else:
-            self.rest_api_key = None
-            self.comet_api = None
+        self._kwargs: Dict[str, Any] = kwargs
 
     @property
     @rank_zero_experiment
@@ -271,50 +261,37 @@ def experiment(self) -> Union["Experiment", "ExistingExperiment", "OfflineExperi
         if self._experiment is not None and self._experiment.alive:
             return self._experiment
 
-        if self._future_experiment_key is not None:
-            os.environ["COMET_EXPERIMENT_KEY"] = self._future_experiment_key
-
-        from comet_ml import ExistingExperiment, Experiment, OfflineExperiment
-
-        try:
-            if self.mode == "online":
-                if self._experiment_key is None:
-                    self._experiment = Experiment(api_key=self.api_key, project_name=self._project_name, **self._kwargs)
-                    self._experiment_key = self._experiment.get_key()
-                else:
-                    self._experiment = ExistingExperiment(
-                        api_key=self.api_key,
-                        project_name=self._project_name,
-                        previous_experiment=self._experiment_key,
-                        **self._kwargs,
-                    )
-            else:
-                self._experiment = OfflineExperiment(
-                    offline_directory=self.save_dir, project_name=self._project_name, **self._kwargs
-                )
-            self._experiment.log_other("Created from", "pytorch-lightning")
-        finally:
-            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)
+        import comet_ml
+
+        comet_comfig = comet_ml.ExperimentConfig(
+            offline_directory=self._save_dir,
+            name=self._experiment_name,
+            **self._kwargs,
+        )
+
+        self._experiment = comet_ml.start(
+            api_key=self.api_key,
+            project=self._project_name,
+            experiment_key=self._experiment_key,
+            online=self.mode == "online",
+            experiment_config=comet_comfig,
+        )
+
+        self._experiment.log_other("Created from", "pytorch-lightning")
 
         return self._experiment
 
     @override
     @rank_zero_only
     def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
         params = _convert_params(params)
-        params = _flatten_dict(params)
         self.experiment.log_parameters(params)
 
     @override
     @rank_zero_only
     def log_metrics(self, metrics: Mapping[str, Union[Tensor, float]], step: Optional[int] = None) -> None:
         assert rank_zero_only.rank == 0, "experiment tried to log from global_rank != 0"
-        # Comet.ml expects metrics to be a dictionary of detached tensors on CPU
+        # Comet.com expects metrics to be a dictionary of detached tensors on CPU
         metrics_without_epoch = metrics.copy()
         for key, val in metrics_without_epoch.items():
             if isinstance(val, Tensor):
@@ -324,26 +301,21 @@ def log_metrics(self, metrics: Mapping[str, Union[Tensor, float]], step: Optiona
         metrics_without_epoch = _add_prefix(metrics_without_epoch, self._prefix, self.LOGGER_JOIN_CHAR)
         self.experiment.log_metrics(metrics_without_epoch, step=step, epoch=epoch)
 
-    def reset_experiment(self) -> None:
-        self._experiment = None
-
     @override
     @rank_zero_only
     def finalize(self, status: str) -> None:
-        r"""When calling ``self.experiment.end()``, that experiment won't log any more data to Comet. That's why, if you
-        need to log any more data, you need to create an ExistingCometExperiment. For example, to log data when testing
-        your model after training, because when training is finalized :meth:`CometLogger.finalize` is called.
-
-        This happens automatically in the :meth:`~CometLogger.experiment` property, when
-        ``self._experiment`` is set to ``None``, i.e. ``self.reset_experiment()``.
+        """
+        We will not end experiment (self._experiment.end()) here
+        to have an ability to continue using it after training is complete
+        but instead of ending we will upload/save all the data
 
         """
         if self._experiment is None:
             # When using multiprocessing, finalize() should be a no-op on the main process, as no experiment has been
             # initialized there
             return
-        self.experiment.end()
-        self.reset_experiment()
+
+        self.experiment.flush()
 
     @property
     @override
@@ -380,35 +352,12 @@ def version(self) -> str:
         """Gets the version.
 
         Returns:
-            The first one of the following that is set in the following order
-
-            1. experiment id.
-            2. experiment key.
-            3. "COMET_EXPERIMENT_KEY" environment variable.
-            4. future experiment key.
-
-            If none are present generates a new guid.
-
+            experiment id/key
         """
-        # Don't create an experiment if we don't have one
         if self._experiment is not None:
             return self._experiment.id
 
-        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
-
-        import comet_ml
-
-        # Pre-generate an experiment key
-        self._future_experiment_key = comet_ml.generate_guid()
-
-        return self._future_experiment_key
+        return self._experiment_key
 
     def __getstate__(self) -> Dict[str, Any]:
         state = self.__dict__.copy()

tests/tests_pytorch/loggers/test_comet.py

@@ -60,11 +60,6 @@ def test_comet_logger_online(comet_mock):
     )
     comet_existing().set_name.assert_called_once_with("experiment")
 
-    # API experiment
-    api = comet_mock.api.API
-    CometLogger(api_key="key", workspace="dummy-test", project_name="general", rest_api_key="rest")
-    api.assert_called_once_with("rest")
-
 
 @mock.patch.dict(os.environ, {})
 def test_comet_experiment_resets_if_not_alive(comet_mock):