Skip to content

feature/comet-logger #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 6 commits into
base: master
Choose a base branch
from
Draft

feature/comet-logger #1

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
dcfc6a5
remove deprecated rest api key
japdubengsub Aug 27, 2024
dcdf053
simplify experiment key configuration/handling
japdubengsub Aug 27, 2024
d810724
use comet_ml.start()
japdubengsub Aug 27, 2024
7538e4e
support nested hyperparams
japdubengsub Aug 27, 2024
009cab0
don't end experiment automatically
japdubengsub Aug 27, 2024
ecc3d98
unify solution with comet_ml.start()
japdubengsub Aug 29, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
141 changes: 45 additions & 96 deletions src/lightning/pytorch/loggers/comet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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):
Expand Down Expand Up @@ -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
)
Expand All @@ -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)
Expand Down Expand Up @@ -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:
Expand All @@ -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,
Expand All @@ -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())

Expand All @@ -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
Expand All @@ -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)
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's good that we don't flatten the dict anymore, did you see a difference in the reported parameters when running the examples?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

for logged parameter:
comet_logger.log_hyperparams({"specific": {'param': {'subparam': "value"}}})

it will look:

name | value
specific | {"param": {"subparam": "value"}}

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):
Expand All @@ -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
Expand Down Expand Up @@ -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()
Expand Down
5 changes: 0 additions & 5 deletions tests/tests_pytorch/loggers/test_comet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
Loading