From ad70d64508deed7d933490eb02b41b5bc28fc4b7 Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Wed, 26 Jun 2024 13:54:15 -0400 Subject: [PATCH 1/7] Adds comet-ml plugin Signed-off-by: Thomas J. Fan --- plugins/flytekit-comet-ml/README.md | 32 ++++ .../flytekitplugins/comet_ml/__init__.py | 3 + .../flytekitplugins/comet_ml/tracking.py | 139 ++++++++++++++++ plugins/flytekit-comet-ml/setup.py | 39 +++++ .../tests/test_comet_ml_init.py | 153 ++++++++++++++++++ 5 files changed, 366 insertions(+) create mode 100644 plugins/flytekit-comet-ml/README.md create mode 100644 plugins/flytekit-comet-ml/flytekitplugins/comet_ml/__init__.py create mode 100644 plugins/flytekit-comet-ml/flytekitplugins/comet_ml/tracking.py create mode 100644 plugins/flytekit-comet-ml/setup.py create mode 100644 plugins/flytekit-comet-ml/tests/test_comet_ml_init.py diff --git a/plugins/flytekit-comet-ml/README.md b/plugins/flytekit-comet-ml/README.md new file mode 100644 index 0000000000..2533c4d4fb --- /dev/null +++ b/plugins/flytekit-comet-ml/README.md @@ -0,0 +1,32 @@ +# Flytekit Comet Plugin + +Comet’s machine learning platform integrates with your existing infrastructure and tools so you can manage, visualize, and optimize models—from training runs to production monitoring. This plugin integrates Flyte with Comet.ml by configuring links between the two platforms. + +To install the plugin, run: + +```bash +pip install flytekitplugins-comet-ml +``` + +Here is an example of running Comet with PyTorch Lightning: + +```python +from flytekit +``` + +Comet requires an API key to authenticate with their platform. In the above example, a secret is created using +[Flyte's Secrets manager](https://docs.flyte.org/en/latest/user_guide/productionizing/secrets.html). + +To enable linking from the Flyte side panel to Comet.ml, add the following to Flyte's configuration: + +```yaml +plugins: + logs: + dynamic-log-links: + - comet-ml-execution-id: + displayName: Comet + templateUris: '{{ .taskConfig.host }}/{{ .taskConfig.entity }}/{{ .taskConfig.project }}/runs/{{ .executionName }}{{ .nodeId }}{{ .taskRetryAttempt }}{{ .taskConfig.experiment_key_suffix }}' + - comet-ml-custom-id: + displayName: Comet + templateUris: '{{ .taskConfig.host }}/{{ .taskConfig.entity }}/{{ .taskConfig.project }}/runs/{{ .taskConfig.id }}' +``` diff --git a/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/__init__.py b/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/__init__.py new file mode 100644 index 0000000000..706e849c34 --- /dev/null +++ b/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/__init__.py @@ -0,0 +1,3 @@ +from .tracking import comet_ml_init + +__all__ = ["comet_ml_init"] diff --git a/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/tracking.py b/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/tracking.py new file mode 100644 index 0000000000..5cc10720b2 --- /dev/null +++ b/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/tracking.py @@ -0,0 +1,139 @@ +import os +from hashlib import shake_256 +from typing import Callable, Optional, Union + +import comet_ml +from flytekit import Secret +from flytekit.core.context_manager import FlyteContextManager +from flytekit.core.utils import ClassDecorator + +COMET_ML_EXECUTION_TYPE_VALUE = "comet-ml-execution-id" +COMET_ML_CUSTOM_TYPE_VALUE = "comet-ml-custom-id" + + +def _generate_suffix_with_length_10(project_name: str, workspace: str) -> str: + """Generate suffix from project_name + workspace.""" + h = shake_256(f"{project_name}-{workspace}".encode("utf-8")) + # Using 5 generates a suffix with length 10 + return h.hexdigest(5) + + +def _generate_experiment_key(hostname: str, project_name: str, workspace: str) -> str: + """Generate experiment key that comet_ml can use: + + 1. Is alphanumeric + 2. 32 <= len(experiment_key) <= 50 + """ + # In Flyte, then hostname is set to {.executionName}-{.nodeID}-{.taskRetryAttempt}, where + # - len(executionName) == 20 + # - 2 <= len(nodeId) <= 8 + # - 1 <= len(taskRetryAttempt)) <= 2 (In practice, retries does not go above 99) + # Removing the `-` because it is not alphanumeric, the 23 <= len(hostname) <= 30 + # On the low end we need to add 10 characters to stay in the range acceptable to comet_ml + hostname = hostname.replace("-", "") + suffix = _generate_suffix_with_length_10(project_name, workspace) + return f"{hostname}{suffix}" + + +class comet_ml_init(ClassDecorator): + COMET_ML_PROJECT_NAME_KEY = "project_name" + COMET_ML_WORKSPACE_KEY = "workspace" + COMET_ML_EXPERIMENT_KEY_KEY = "experiment_key" + COMET_ML_URL_SUFFIX_KEY = "link_suffix" + + def __init__( + self, + task_function: Optional[Callable] = None, + project_name: Optional[str] = None, + workspace: Optional[str] = None, + experiment_key: Optional[str] = None, + secret: Optional[Union[Secret, Callable]] = None, + **init_kwargs: dict, + ): + """Comet plugin. + Args: + task_function (function, optional): The user function to be decorated. Defaults to None. + project_name (str): Send your experiment to a specific project. (Required) + workspace (str): Attach an experiment to a project that belongs to this workspace. (Required) + experiment_key (str): Experiment key. + secret (Secret or Callable): Secret with your `COMET_API_KEY` or a callable that returns the API key. + The callable takes no arguments and returns a string. (Required) + **init_kwargs (dict): The rest of the arguments are passed directly to `comet_ml.init`. + """ + if project_name is None: + raise ValueError("project_name must be set") + if workspace is None: + raise ValueError("workspace must be set") + if secret is None: + raise ValueError("secret must be set") + + self.project_name = project_name + self.workspace = workspace + self.experiment_key = experiment_key + self.secret = secret + self.init_kwargs = init_kwargs + + super().__init__( + task_function, + project_name=project_name, + workspace=workspace, + experiment_key=experiment_key, + secret=secret, + **init_kwargs, + ) + + def execute(self, *args, **kwargs): + ctx = FlyteContextManager.current_context() + is_local_execution = ctx.execution_state.is_local_execution() + + default_kwargs = self.init_kwargs + init_kwargs = { + "project_name": self.project_name, + "workspace": self.workspace, + **default_kwargs, + } + + if is_local_execution: + # For local execution, always use the experiment_key. If `self.experiment_key` is `None`, comet_ml + # will generate it's own key + init_kwargs["experiment_key"] = self.experiment_key + else: + # Get api key for remote execution + if isinstance(self.secret, Secret): + secrets = ctx.user_space_params.secrets + comet_ml_api_key = secrets.get(key=self.secret.key, group=self.secret.group) + else: + comet_ml_api_key = self.secret() + + init_kwargs["api_key"] = comet_ml_api_key + + if self.experiment_key is None: + # The HOSTNAME is set to {.executionName}-{.nodeID}-{.taskRetryAttempt} + # If HOSTNAME is not defined, use the execution name as a fallback + hostname = os.environ.get("HOSTNAME", ctx.user_space_params.execution_id.name) + experiment_key = _generate_experiment_key(hostname, self.project_name, self.workspace) + else: + experiment_key = self.experiment_key + + init_kwargs["experiment_key"] = experiment_key + + comet_ml.init(**init_kwargs) + output = self.task_function(*args, **kwargs) + return output + + def get_extra_config(self): + extra_config = { + self.COMET_ML_PROJECT_NAME_KEY: self.project_name, + self.COMET_ML_WORKSPACE_KEY: self.workspace, + } + + if self.experiment_key is None: + comet_ml_value = COMET_ML_EXECUTION_TYPE_VALUE + suffix = _generate_suffix_with_length_10(self.project_name, self.workspace) + extra_config[self.COMET_ML_URL_SUFFIX_KEY] = suffix + else: + comet_ml_value = COMET_ML_CUSTOM_TYPE_VALUE + extra_config[self.COMET_ML_EXPERIMENT_KEY_KEY] = self.experiment_key + + extra_config[self.LINK_TYPE_KEY] = comet_ml_value + return extra_config diff --git a/plugins/flytekit-comet-ml/setup.py b/plugins/flytekit-comet-ml/setup.py new file mode 100644 index 0000000000..387b9119e3 --- /dev/null +++ b/plugins/flytekit-comet-ml/setup.py @@ -0,0 +1,39 @@ +from setuptools import setup + +PLUGIN_NAME = "comet-ml" +MODULE_NAME = "comet_ml" + + +microlib_name = f"flytekitplugins-{PLUGIN_NAME}" + +plugin_requires = ["flytekit>=1.12.3", "comet-ml>=3.43.2"] + +__version__ = "0.0.0+develop" + +setup( + name=microlib_name, + version=__version__, + author="flyteorg", + author_email="admin@flyte.org", + description="This package enables seamless use of Comet within Flyte", + namespace_packages=["flytekitplugins"], + packages=[f"flytekitplugins.{MODULE_NAME}"], + install_requires=plugin_requires, + license="apache2", + python_requires=">=3.8", + classifiers=[ + "Intended Audience :: Science/Research", + "Intended Audience :: Developers", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Topic :: Scientific/Engineering", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Topic :: Software Development", + "Topic :: Software Development :: Libraries", + "Topic :: Software Development :: Libraries :: Python Modules", + ], +) diff --git a/plugins/flytekit-comet-ml/tests/test_comet_ml_init.py b/plugins/flytekit-comet-ml/tests/test_comet_ml_init.py new file mode 100644 index 0000000000..9216106428 --- /dev/null +++ b/plugins/flytekit-comet-ml/tests/test_comet_ml_init.py @@ -0,0 +1,153 @@ +from hashlib import shake_256 +from unittest.mock import patch, Mock +import pytest + +from flytekit import Secret, task +from flytekitplugins.comet_ml import comet_ml_init +from flytekitplugins.comet_ml.tracking import COMET_ML_CUSTOM_TYPE_VALUE, COMET_ML_EXECUTION_TYPE_VALUE, _generate_suffix_with_length_10, _generate_experiment_key + + +secret = Secret(key="abc", group="xyz") + + +@pytest.mark.parametrize("experiment_key", [None, "abc123dfassfasfsafsafd"]) +def test_extra_config(experiment_key): + project_name = "abc" + workspace = "my_workspace" + + comet_decorator = comet_ml_init( + project_name=project_name, + workspace=workspace, + experiment_key=experiment_key, + secret=secret + ) + + assert comet_decorator.secret is secret + extra_config = comet_decorator.get_extra_config() + + if experiment_key is None: + assert extra_config[comet_decorator.LINK_TYPE_KEY] == COMET_ML_EXECUTION_TYPE_VALUE + assert comet_decorator.COMET_ML_EXPERIMENT_KEY_KEY not in extra_config + + suffix = _generate_suffix_with_length_10(project_name=project_name, workspace=workspace) + assert extra_config[comet_decorator.COMET_ML_URL_SUFFIX_KEY] == suffix + + else: + assert extra_config[comet_decorator.LINK_TYPE_KEY] == COMET_ML_CUSTOM_TYPE_VALUE + assert extra_config[comet_decorator.COMET_ML_EXPERIMENT_KEY_KEY] == experiment_key + assert comet_decorator.COMET_ML_URL_SUFFIX_KEY not in extra_config + + assert extra_config[comet_decorator.COMET_ML_WORKSPACE_KEY] == workspace + + +@task +@comet_ml_init(project_name="abc", workspace="my-workspace", secret=secret, log_code=False) +def train_model(): + pass + + +@patch("flytekitplugins.comet_ml.tracking.comet_ml") +def test_local_execution(comet_ml_mock): + train_model() + + comet_ml_mock.init.assert_called_with(project_name="abc", workspace="my-workspace", log_code=False, experiment_key=None) + + +@task +@comet_ml_init( + project_name="xyz", + workspace="another-workspace", + secret=secret, + experiment_key="my-previous-experiment-key", +) +def train_model_with_experiment_key(): + pass + + +@patch("flytekitplugins.comet_ml.tracking.comet_ml") +def test_local_execution_with_experiment_key(comet_ml_mock): + train_model_with_experiment_key() + + comet_ml_mock.init.assert_called_with( + project_name="xyz", + workspace="another-workspace", + experiment_key="my-previous-experiment-key", + ) + + +@patch("flytekitplugins.comet_ml.tracking.os") +@patch("flytekitplugins.comet_ml.tracking.FlyteContextManager") +@patch("flytekitplugins.comet_ml.tracking.comet_ml") +def test_remote_execution(comet_ml_mock, manager_mock, os_mock): + # Pretend that the execution is remote + ctx_mock = Mock() + ctx_mock.execution_state.is_local_execution.return_value = False + + ctx_mock.user_space_params.secrets.get.return_value = "this_is_the_secret" + ctx_mock.user_space_params.execution_id.name = "my_execution_id" + + manager_mock.current_context.return_value = ctx_mock + hostname = "a423423423afasf4jigl-fasj4321-0" + os_mock.environ = {"HOSTNAME": hostname} + + project_name = "abc" + workspace = "my-workspace" + + h = shake_256(f"{project_name}-{workspace}".encode("utf-8")) + suffix = h.hexdigest(5) + hostname_alpha = hostname.replace("-", "") + experiment_key = f"{hostname_alpha}{suffix}" + + train_model() + + comet_ml_mock.init.assert_called_with( + project_name="abc", + workspace="my-workspace", + api_key="this_is_the_secret", + experiment_key=experiment_key, + log_code=False, + ) + ctx_mock.user_space_params.secrets.get.assert_called_with(key="abc", group="xyz") + + +def get_secret(): + return "my-comet-ml-api-key" + + +@task +@comet_ml_init(project_name="my_project", workspace="my_workspace", secret=get_secret) +def train_model_with_callable_secret(): + pass + + +@patch("flytekitplugins.comet_ml.tracking.os") +@patch("flytekitplugins.comet_ml.tracking.FlyteContextManager") +@patch("flytekitplugins.comet_ml.tracking.comet_ml") +def test_remote_execution_with_callable_secret(comet_ml_mock, manager_mock, os_mock): + # Pretend that the execution is remote + ctx_mock = Mock() + ctx_mock.execution_state.is_local_execution.return_value = False + + manager_mock.current_context.return_value = ctx_mock + hostname = "a423423423afasf4jigl-fasj4321-0" + os_mock.environ = {"HOSTNAME": hostname} + + train_model_with_callable_secret() + + comet_ml_mock.init.assert_called_with( + project_name="my_project", + api_key="my-comet-ml-api-key", + workspace="my_workspace", + experiment_key=_generate_experiment_key(hostname, "my_project", "my_workspace") + ) + + +def test_errors(): + with pytest.raises(ValueError, match="project_name must be set"): + comet_ml_init() + + with pytest.raises(ValueError, match="workspace must be set"): + comet_ml_init(project_name="abc") + + with pytest.raises(ValueError, match="secret must be set"): + comet_ml_init(project_name="abc", workspace="xyz") From da94ee269535b697e39dd55c202ea414d56a44f6 Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Thu, 27 Jun 2024 10:51:00 -0400 Subject: [PATCH 2/7] For local execution, do not set experiment_key if it is none Signed-off-by: Thomas J. Fan --- .../flytekitplugins/comet_ml/tracking.py | 3 ++- plugins/flytekit-comet-ml/tests/test_comet_ml_init.py | 10 ++++++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/tracking.py b/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/tracking.py index 5cc10720b2..32a23505ad 100644 --- a/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/tracking.py +++ b/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/tracking.py @@ -96,7 +96,8 @@ def execute(self, *args, **kwargs): if is_local_execution: # For local execution, always use the experiment_key. If `self.experiment_key` is `None`, comet_ml # will generate it's own key - init_kwargs["experiment_key"] = self.experiment_key + if self.experiment_key is not None: + init_kwargs["experiment_key"] = self.experiment_key else: # Get api key for remote execution if isinstance(self.secret, Secret): diff --git a/plugins/flytekit-comet-ml/tests/test_comet_ml_init.py b/plugins/flytekit-comet-ml/tests/test_comet_ml_init.py index 9216106428..c12a0eab15 100644 --- a/plugins/flytekit-comet-ml/tests/test_comet_ml_init.py +++ b/plugins/flytekit-comet-ml/tests/test_comet_ml_init.py @@ -4,7 +4,12 @@ from flytekit import Secret, task from flytekitplugins.comet_ml import comet_ml_init -from flytekitplugins.comet_ml.tracking import COMET_ML_CUSTOM_TYPE_VALUE, COMET_ML_EXECUTION_TYPE_VALUE, _generate_suffix_with_length_10, _generate_experiment_key +from flytekitplugins.comet_ml.tracking import ( + COMET_ML_CUSTOM_TYPE_VALUE, + COMET_ML_EXECUTION_TYPE_VALUE, + _generate_suffix_with_length_10, + _generate_experiment_key, +) secret = Secret(key="abc", group="xyz") @@ -50,7 +55,8 @@ def train_model(): def test_local_execution(comet_ml_mock): train_model() - comet_ml_mock.init.assert_called_with(project_name="abc", workspace="my-workspace", log_code=False, experiment_key=None) + comet_ml_mock.init.assert_called_with( + project_name="abc", workspace="my-workspace", log_code=False) @task From 117f05063512263df4357185144380187c96cd7f Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Thu, 27 Jun 2024 23:11:49 -0400 Subject: [PATCH 3/7] Use correct comet-ml links Signed-off-by: Thomas J. Fan --- plugins/flytekit-comet-ml/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/plugins/flytekit-comet-ml/README.md b/plugins/flytekit-comet-ml/README.md index 2533c4d4fb..3bfe73c10b 100644 --- a/plugins/flytekit-comet-ml/README.md +++ b/plugins/flytekit-comet-ml/README.md @@ -25,8 +25,8 @@ plugins: dynamic-log-links: - comet-ml-execution-id: displayName: Comet - templateUris: '{{ .taskConfig.host }}/{{ .taskConfig.entity }}/{{ .taskConfig.project }}/runs/{{ .executionName }}{{ .nodeId }}{{ .taskRetryAttempt }}{{ .taskConfig.experiment_key_suffix }}' + templateUris: https://www.comet.com/{{ .taskConfig.workspace }}/{{ .taskConfig.project_name }}/{{ .executionName }}{{ .nodeId }}{{ .taskRetryAttempt }}{{ .taskConfig.link_suffix }} - comet-ml-custom-id: displayName: Comet - templateUris: '{{ .taskConfig.host }}/{{ .taskConfig.entity }}/{{ .taskConfig.project }}/runs/{{ .taskConfig.id }}' + templateUris: https://www.comet.com/{{ .taskConfig.workspace }}/{{ .taskConfig.project_name }}/{{ .taskConfig.experiment_key }} ``` From 4cfa4eb35d1997d81fd0d86f4f4860844529ce5a Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Tue, 2 Jul 2024 13:04:59 -0400 Subject: [PATCH 4/7] Allow host to be adjustable Signed-off-by: Thomas J. Fan --- plugins/flytekit-comet-ml/README.md | 4 ++-- .../flytekit-comet-ml/flytekitplugins/comet_ml/tracking.py | 6 ++++++ plugins/flytekit-comet-ml/tests/test_comet_ml_init.py | 1 + 3 files changed, 9 insertions(+), 2 deletions(-) diff --git a/plugins/flytekit-comet-ml/README.md b/plugins/flytekit-comet-ml/README.md index 3bfe73c10b..99e475e97a 100644 --- a/plugins/flytekit-comet-ml/README.md +++ b/plugins/flytekit-comet-ml/README.md @@ -25,8 +25,8 @@ plugins: dynamic-log-links: - comet-ml-execution-id: displayName: Comet - templateUris: https://www.comet.com/{{ .taskConfig.workspace }}/{{ .taskConfig.project_name }}/{{ .executionName }}{{ .nodeId }}{{ .taskRetryAttempt }}{{ .taskConfig.link_suffix }} + templateUris: {{ .taskConfig.host }}/{{ .taskConfig.workspace }}/{{ .taskConfig.project_name }}/{{ .executionName }}{{ .nodeId }}{{ .taskRetryAttempt }}{{ .taskConfig.link_suffix }} - comet-ml-custom-id: displayName: Comet - templateUris: https://www.comet.com/{{ .taskConfig.workspace }}/{{ .taskConfig.project_name }}/{{ .taskConfig.experiment_key }} + templateUris: {{ .taskConfig.host }}/{{ .taskConfig.workspace }}/{{ .taskConfig.project_name }}/{{ .taskConfig.experiment_key }} ``` diff --git a/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/tracking.py b/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/tracking.py index 32a23505ad..ae44b1ae04 100644 --- a/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/tracking.py +++ b/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/tracking.py @@ -40,6 +40,7 @@ class comet_ml_init(ClassDecorator): COMET_ML_WORKSPACE_KEY = "workspace" COMET_ML_EXPERIMENT_KEY_KEY = "experiment_key" COMET_ML_URL_SUFFIX_KEY = "link_suffix" + COMET_ML_HOST_KEY = "host" def __init__( self, @@ -48,6 +49,7 @@ def __init__( workspace: Optional[str] = None, experiment_key: Optional[str] = None, secret: Optional[Union[Secret, Callable]] = None, + host: str = "https://www.comet.com", **init_kwargs: dict, ): """Comet plugin. @@ -58,6 +60,7 @@ def __init__( experiment_key (str): Experiment key. secret (Secret or Callable): Secret with your `COMET_API_KEY` or a callable that returns the API key. The callable takes no arguments and returns a string. (Required) + host (str): URL to your Comet service. Defaults to "https://www.comet.com" **init_kwargs (dict): The rest of the arguments are passed directly to `comet_ml.init`. """ if project_name is None: @@ -71,6 +74,7 @@ def __init__( self.workspace = workspace self.experiment_key = experiment_key self.secret = secret + self.host = host self.init_kwargs = init_kwargs super().__init__( @@ -79,6 +83,7 @@ def __init__( workspace=workspace, experiment_key=experiment_key, secret=secret, + host=host, **init_kwargs, ) @@ -126,6 +131,7 @@ def get_extra_config(self): extra_config = { self.COMET_ML_PROJECT_NAME_KEY: self.project_name, self.COMET_ML_WORKSPACE_KEY: self.workspace, + self.COMET_ML_HOST_KEY: self.host, } if self.experiment_key is None: diff --git a/plugins/flytekit-comet-ml/tests/test_comet_ml_init.py b/plugins/flytekit-comet-ml/tests/test_comet_ml_init.py index c12a0eab15..fe204d357a 100644 --- a/plugins/flytekit-comet-ml/tests/test_comet_ml_init.py +++ b/plugins/flytekit-comet-ml/tests/test_comet_ml_init.py @@ -43,6 +43,7 @@ def test_extra_config(experiment_key): assert comet_decorator.COMET_ML_URL_SUFFIX_KEY not in extra_config assert extra_config[comet_decorator.COMET_ML_WORKSPACE_KEY] == workspace + assert extra_config[comet_decorator.COMET_ML_HOST_KEY] == "https://www.comet.com" @task From 58b33b0b4d1eff5be5f3c76a04c299b36ecbb406 Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Tue, 2 Jul 2024 22:34:46 -0400 Subject: [PATCH 5/7] Adds comet-ml plugin Signed-off-by: Thomas J. Fan --- plugins/flytekit-comet-ml/README.md | 6 ------ 1 file changed, 6 deletions(-) diff --git a/plugins/flytekit-comet-ml/README.md b/plugins/flytekit-comet-ml/README.md index 99e475e97a..3d90924445 100644 --- a/plugins/flytekit-comet-ml/README.md +++ b/plugins/flytekit-comet-ml/README.md @@ -8,12 +8,6 @@ To install the plugin, run: pip install flytekitplugins-comet-ml ``` -Here is an example of running Comet with PyTorch Lightning: - -```python -from flytekit -``` - Comet requires an API key to authenticate with their platform. In the above example, a secret is created using [Flyte's Secrets manager](https://docs.flyte.org/en/latest/user_guide/productionizing/secrets.html). From d04d0129f99eb7e7edcffeae04d4f401f49be6d6 Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Wed, 17 Jul 2024 12:44:30 -0400 Subject: [PATCH 6/7] Use new comet-ml login name Signed-off-by: Thomas J. Fan --- .../flytekitplugins/comet_ml/__init__.py | 4 +-- .../flytekitplugins/comet_ml/tracking.py | 26 +++++++++++-------- .../tests/test_comet_ml_init.py | 16 ++++++------ 3 files changed, 25 insertions(+), 21 deletions(-) diff --git a/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/__init__.py b/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/__init__.py index 706e849c34..58dbff81d2 100644 --- a/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/__init__.py +++ b/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/__init__.py @@ -1,3 +1,3 @@ -from .tracking import comet_ml_init +from .tracking import comet_ml_login -__all__ = ["comet_ml_init"] +__all__ = ["comet_ml_login"] diff --git a/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/tracking.py b/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/tracking.py index ae44b1ae04..77ec31a179 100644 --- a/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/tracking.py +++ b/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/tracking.py @@ -35,7 +35,7 @@ def _generate_experiment_key(hostname: str, project_name: str, workspace: str) - return f"{hostname}{suffix}" -class comet_ml_init(ClassDecorator): +class comet_ml_login(ClassDecorator): COMET_ML_PROJECT_NAME_KEY = "project_name" COMET_ML_WORKSPACE_KEY = "workspace" COMET_ML_EXPERIMENT_KEY_KEY = "experiment_key" @@ -50,7 +50,7 @@ def __init__( experiment_key: Optional[str] = None, secret: Optional[Union[Secret, Callable]] = None, host: str = "https://www.comet.com", - **init_kwargs: dict, + **login_kwargs: dict, ): """Comet plugin. Args: @@ -61,7 +61,7 @@ def __init__( secret (Secret or Callable): Secret with your `COMET_API_KEY` or a callable that returns the API key. The callable takes no arguments and returns a string. (Required) host (str): URL to your Comet service. Defaults to "https://www.comet.com" - **init_kwargs (dict): The rest of the arguments are passed directly to `comet_ml.init`. + **login_kwargs (dict): The rest of the arguments are passed directly to `comet_ml.login`. """ if project_name is None: raise ValueError("project_name must be set") @@ -75,7 +75,7 @@ def __init__( self.experiment_key = experiment_key self.secret = secret self.host = host - self.init_kwargs = init_kwargs + self.login_kwargs = login_kwargs super().__init__( task_function, @@ -84,15 +84,15 @@ def __init__( experiment_key=experiment_key, secret=secret, host=host, - **init_kwargs, + **login_kwargs, ) def execute(self, *args, **kwargs): ctx = FlyteContextManager.current_context() is_local_execution = ctx.execution_state.is_local_execution() - default_kwargs = self.init_kwargs - init_kwargs = { + default_kwargs = self.login_kwargs + login_kwargs = { "project_name": self.project_name, "workspace": self.workspace, **default_kwargs, @@ -102,7 +102,7 @@ def execute(self, *args, **kwargs): # For local execution, always use the experiment_key. If `self.experiment_key` is `None`, comet_ml # will generate it's own key if self.experiment_key is not None: - init_kwargs["experiment_key"] = self.experiment_key + login_kwargs["experiment_key"] = self.experiment_key else: # Get api key for remote execution if isinstance(self.secret, Secret): @@ -111,7 +111,7 @@ def execute(self, *args, **kwargs): else: comet_ml_api_key = self.secret() - init_kwargs["api_key"] = comet_ml_api_key + login_kwargs["api_key"] = comet_ml_api_key if self.experiment_key is None: # The HOSTNAME is set to {.executionName}-{.nodeID}-{.taskRetryAttempt} @@ -121,9 +121,13 @@ def execute(self, *args, **kwargs): else: experiment_key = self.experiment_key - init_kwargs["experiment_key"] = experiment_key + login_kwargs["experiment_key"] = experiment_key + + if hasattr(comet_ml, "login"): + comet_ml.login(**login_kwargs) + else: + comet_ml.init(**login_kwargs) - comet_ml.init(**init_kwargs) output = self.task_function(*args, **kwargs) return output diff --git a/plugins/flytekit-comet-ml/tests/test_comet_ml_init.py b/plugins/flytekit-comet-ml/tests/test_comet_ml_init.py index fe204d357a..28f5b56b86 100644 --- a/plugins/flytekit-comet-ml/tests/test_comet_ml_init.py +++ b/plugins/flytekit-comet-ml/tests/test_comet_ml_init.py @@ -3,7 +3,7 @@ import pytest from flytekit import Secret, task -from flytekitplugins.comet_ml import comet_ml_init +from flytekitplugins.comet_ml import comet_ml_login from flytekitplugins.comet_ml.tracking import ( COMET_ML_CUSTOM_TYPE_VALUE, COMET_ML_EXECUTION_TYPE_VALUE, @@ -20,7 +20,7 @@ def test_extra_config(experiment_key): project_name = "abc" workspace = "my_workspace" - comet_decorator = comet_ml_init( + comet_decorator = comet_ml_login( project_name=project_name, workspace=workspace, experiment_key=experiment_key, @@ -47,7 +47,7 @@ def test_extra_config(experiment_key): @task -@comet_ml_init(project_name="abc", workspace="my-workspace", secret=secret, log_code=False) +@comet_ml_login(project_name="abc", workspace="my-workspace", secret=secret, log_code=False) def train_model(): pass @@ -61,7 +61,7 @@ def test_local_execution(comet_ml_mock): @task -@comet_ml_init( +@comet_ml_login( project_name="xyz", workspace="another-workspace", secret=secret, @@ -122,7 +122,7 @@ def get_secret(): @task -@comet_ml_init(project_name="my_project", workspace="my_workspace", secret=get_secret) +@comet_ml_login(project_name="my_project", workspace="my_workspace", secret=get_secret) def train_model_with_callable_secret(): pass @@ -151,10 +151,10 @@ def test_remote_execution_with_callable_secret(comet_ml_mock, manager_mock, os_m def test_errors(): with pytest.raises(ValueError, match="project_name must be set"): - comet_ml_init() + comet_ml_login() with pytest.raises(ValueError, match="workspace must be set"): - comet_ml_init(project_name="abc") + comet_ml_login(project_name="abc") with pytest.raises(ValueError, match="secret must be set"): - comet_ml_init(project_name="abc", workspace="xyz") + comet_ml_login(project_name="abc", workspace="xyz") From cc2e9ad8788f7dc846fb5b798e059c21b59837be Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Wed, 17 Jul 2024 15:21:42 -0400 Subject: [PATCH 7/7] Require the project_name workspace and secrets Signed-off-by: Thomas J. Fan --- .github/workflows/pythonbuild.yml | 1 + plugins/flytekit-comet-ml/README.md | 4 +- .../flytekitplugins/comet_ml/tracking.py | 49 ++++++++++++++----- .../tests/test_comet_ml_init.py | 43 +++++++--------- 4 files changed, 57 insertions(+), 40 deletions(-) diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index 8d450ffc87..005658497b 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -316,6 +316,7 @@ jobs: - flytekit-aws-batch - flytekit-aws-sagemaker - flytekit-bigquery + - flytekit-comet-ml - flytekit-dask - flytekit-data-fsspec - flytekit-dbt diff --git a/plugins/flytekit-comet-ml/README.md b/plugins/flytekit-comet-ml/README.md index 3d90924445..a7038c8caf 100644 --- a/plugins/flytekit-comet-ml/README.md +++ b/plugins/flytekit-comet-ml/README.md @@ -19,8 +19,8 @@ plugins: dynamic-log-links: - comet-ml-execution-id: displayName: Comet - templateUris: {{ .taskConfig.host }}/{{ .taskConfig.workspace }}/{{ .taskConfig.project_name }}/{{ .executionName }}{{ .nodeId }}{{ .taskRetryAttempt }}{{ .taskConfig.link_suffix }} + templateUris: "{{ .taskConfig.host }}/{{ .taskConfig.workspace }}/{{ .taskConfig.project_name }}/{{ .executionName }}{{ .nodeId }}{{ .taskRetryAttempt }}{{ .taskConfig.link_suffix }}" - comet-ml-custom-id: displayName: Comet - templateUris: {{ .taskConfig.host }}/{{ .taskConfig.workspace }}/{{ .taskConfig.project_name }}/{{ .taskConfig.experiment_key }} + templateUris: "{{ .taskConfig.host }}/{{ .taskConfig.workspace }}/{{ .taskConfig.project_name }}/{{ .taskConfig.experiment_key }}" ``` diff --git a/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/tracking.py b/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/tracking.py index 77ec31a179..3014513d0d 100644 --- a/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/tracking.py +++ b/plugins/flytekit-comet-ml/flytekitplugins/comet_ml/tracking.py @@ -1,4 +1,5 @@ import os +from functools import partial from hashlib import shake_256 from typing import Callable, Optional, Union @@ -35,7 +36,36 @@ def _generate_experiment_key(hostname: str, project_name: str, workspace: str) - return f"{hostname}{suffix}" -class comet_ml_login(ClassDecorator): +def comet_ml_login( + project_name: str, + workspace: str, + secret: Union[Secret, Callable], + experiment_key: Optional[str] = None, + host: str = "https://www.comet.com", + **login_kwargs: dict, +): + """Comet plugin. + Args: + project_name (str): Send your experiment to a specific project. (Required) + workspace (str): Attach an experiment to a project that belongs to this workspace. (Required) + secret (Secret or Callable): Secret with your `COMET_API_KEY` or a callable that returns the API key. + The callable takes no arguments and returns a string. (Required) + experiment_key (str): Experiment key. + host (str): URL to your Comet service. Defaults to "https://www.comet.com" + **login_kwargs (dict): The rest of the arguments are passed directly to `comet_ml.login`. + """ + return partial( + _comet_ml_login_class, + project_name=project_name, + workspace=workspace, + secret=secret, + experiment_key=experiment_key, + host=host, + **login_kwargs, + ) + + +class _comet_ml_login_class(ClassDecorator): COMET_ML_PROJECT_NAME_KEY = "project_name" COMET_ML_WORKSPACE_KEY = "workspace" COMET_ML_EXPERIMENT_KEY_KEY = "experiment_key" @@ -44,31 +74,24 @@ class comet_ml_login(ClassDecorator): def __init__( self, - task_function: Optional[Callable] = None, - project_name: Optional[str] = None, - workspace: Optional[str] = None, + task_function: Callable, + project_name: str, + workspace: str, + secret: Union[Secret, Callable], experiment_key: Optional[str] = None, - secret: Optional[Union[Secret, Callable]] = None, host: str = "https://www.comet.com", **login_kwargs: dict, ): """Comet plugin. Args: - task_function (function, optional): The user function to be decorated. Defaults to None. project_name (str): Send your experiment to a specific project. (Required) workspace (str): Attach an experiment to a project that belongs to this workspace. (Required) - experiment_key (str): Experiment key. secret (Secret or Callable): Secret with your `COMET_API_KEY` or a callable that returns the API key. The callable takes no arguments and returns a string. (Required) + experiment_key (str): Experiment key. host (str): URL to your Comet service. Defaults to "https://www.comet.com" **login_kwargs (dict): The rest of the arguments are passed directly to `comet_ml.login`. """ - if project_name is None: - raise ValueError("project_name must be set") - if workspace is None: - raise ValueError("workspace must be set") - if secret is None: - raise ValueError("secret must be set") self.project_name = project_name self.workspace = workspace diff --git a/plugins/flytekit-comet-ml/tests/test_comet_ml_init.py b/plugins/flytekit-comet-ml/tests/test_comet_ml_init.py index 28f5b56b86..5572e4a56e 100644 --- a/plugins/flytekit-comet-ml/tests/test_comet_ml_init.py +++ b/plugins/flytekit-comet-ml/tests/test_comet_ml_init.py @@ -27,23 +27,27 @@ def test_extra_config(experiment_key): secret=secret ) - assert comet_decorator.secret is secret - extra_config = comet_decorator.get_extra_config() + @comet_decorator + def task(): + pass + + assert task.secret is secret + extra_config = task.get_extra_config() if experiment_key is None: - assert extra_config[comet_decorator.LINK_TYPE_KEY] == COMET_ML_EXECUTION_TYPE_VALUE - assert comet_decorator.COMET_ML_EXPERIMENT_KEY_KEY not in extra_config + assert extra_config[task.LINK_TYPE_KEY] == COMET_ML_EXECUTION_TYPE_VALUE + assert task.COMET_ML_EXPERIMENT_KEY_KEY not in extra_config suffix = _generate_suffix_with_length_10(project_name=project_name, workspace=workspace) - assert extra_config[comet_decorator.COMET_ML_URL_SUFFIX_KEY] == suffix + assert extra_config[task.COMET_ML_URL_SUFFIX_KEY] == suffix else: - assert extra_config[comet_decorator.LINK_TYPE_KEY] == COMET_ML_CUSTOM_TYPE_VALUE - assert extra_config[comet_decorator.COMET_ML_EXPERIMENT_KEY_KEY] == experiment_key - assert comet_decorator.COMET_ML_URL_SUFFIX_KEY not in extra_config + assert extra_config[task.LINK_TYPE_KEY] == COMET_ML_CUSTOM_TYPE_VALUE + assert extra_config[task.COMET_ML_EXPERIMENT_KEY_KEY] == experiment_key + assert task.COMET_ML_URL_SUFFIX_KEY not in extra_config - assert extra_config[comet_decorator.COMET_ML_WORKSPACE_KEY] == workspace - assert extra_config[comet_decorator.COMET_ML_HOST_KEY] == "https://www.comet.com" + assert extra_config[task.COMET_ML_WORKSPACE_KEY] == workspace + assert extra_config[task.COMET_ML_HOST_KEY] == "https://www.comet.com" @task @@ -56,7 +60,7 @@ def train_model(): def test_local_execution(comet_ml_mock): train_model() - comet_ml_mock.init.assert_called_with( + comet_ml_mock.login.assert_called_with( project_name="abc", workspace="my-workspace", log_code=False) @@ -75,7 +79,7 @@ def train_model_with_experiment_key(): def test_local_execution_with_experiment_key(comet_ml_mock): train_model_with_experiment_key() - comet_ml_mock.init.assert_called_with( + comet_ml_mock.login.assert_called_with( project_name="xyz", workspace="another-workspace", experiment_key="my-previous-experiment-key", @@ -107,7 +111,7 @@ def test_remote_execution(comet_ml_mock, manager_mock, os_mock): train_model() - comet_ml_mock.init.assert_called_with( + comet_ml_mock.login.assert_called_with( project_name="abc", workspace="my-workspace", api_key="this_is_the_secret", @@ -141,20 +145,9 @@ def test_remote_execution_with_callable_secret(comet_ml_mock, manager_mock, os_m train_model_with_callable_secret() - comet_ml_mock.init.assert_called_with( + comet_ml_mock.login.assert_called_with( project_name="my_project", api_key="my-comet-ml-api-key", workspace="my_workspace", experiment_key=_generate_experiment_key(hostname, "my_project", "my_workspace") ) - - -def test_errors(): - with pytest.raises(ValueError, match="project_name must be set"): - comet_ml_login() - - with pytest.raises(ValueError, match="workspace must be set"): - comet_ml_login(project_name="abc") - - with pytest.raises(ValueError, match="secret must be set"): - comet_ml_login(project_name="abc", workspace="xyz")