Require the project_name workspace and secrets

Signed-off-by: Thomas J. Fan <[email protected]>

.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

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 }}"
 ```

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

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")