From cc2e9ad8788f7dc846fb5b798e059c21b59837be Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Wed, 17 Jul 2024 15:21:42 -0400 Subject: [PATCH] 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")