Skip to content

Commit

Permalink
Reapply "Require the project_name workspace and secrets"
Browse files Browse the repository at this point in the history
This reverts commit 216359e.
  • Loading branch information
thomasjpfan committed Jul 17, 2024
1 parent 216359e commit 417d03b
Show file tree
Hide file tree
Showing 4 changed files with 57 additions and 40 deletions.
1 change: 1 addition & 0 deletions .github/workflows/pythonbuild.yml
Original file line number Diff line number Diff line change
Expand Up @@ -316,6 +316,7 @@ jobs:
- flytekit-aws-batch
- flytekit-aws-sagemaker
- flytekit-bigquery
- flytekit-comet-ml
- flytekit-dask
- flytekit-data-fsspec
- flytekit-dbt
Expand Down
4 changes: 2 additions & 2 deletions plugins/flytekit-comet-ml/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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 }}"
```
49 changes: 36 additions & 13 deletions plugins/flytekit-comet-ml/flytekitplugins/comet_ml/tracking.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import os
from functools import partial
from hashlib import shake_256
from typing import Callable, Optional, Union

Expand Down Expand Up @@ -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"
Expand All @@ -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
Expand Down
43 changes: 18 additions & 25 deletions plugins/flytekit-comet-ml/tests/test_comet_ml_init.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)


Expand All @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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")

0 comments on commit 417d03b

Please sign in to comment.