Skip to content
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

Jump to bottom

Add optional experiment_id parameter to mlflow.set_experiment #5012

Merged
merged 5 commits into from Nov 5, 2021
Merged

Add optional experiment_id parameter to mlflow.set_experiment #5012

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
11283dc
Impl
dbczumar Nov 5, 2021
bb3ee5b
PR
dbczumar Nov 5, 2021
614a0f8
Format
dbczumar Nov 5, 2021
fee73f9
Address comments
dbczumar Nov 5, 2021
34ad0a0
Return experiment
dbczumar Nov 5, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
33 changes: 16 additions & 17 deletions mlflow/tracking/fluent.py
View file
Open in desktop
Expand Up @@ -65,6 +65,8 @@ def set_experiment(experiment_name: str = None, experiment_id: str = None) -> No
created.
:param experiment_id: ID of the experiment to be activated. If an experiment with this ID
does not exist, an exception is thrown.
:return: An instance of :py:class:`mlflow.entities.Experiment` representing the new active
experiment.

.. code-block:: python
:caption: Example
Expand Down Expand Up @@ -97,24 +99,10 @@ def set_experiment(experiment_name: str = None, experiment_id: str = None) -> No
error_code=INVALID_PARAMETER_VALUE,
)

def verify_experiment_active(experiment):
if experiment.lifecycle_stage != LifecycleStage.ACTIVE:
raise MlflowException(
message=(
"Cannot set a deleted experiment '%s' as the active experiment."
" You can restore the experiment, or permanently delete the "
" experiment to create a new one." % experiment.name
),
error_code=INVALID_PARAMETER_VALUE,
)

client = MlflowClient()
if experiment_id is None:
experiment = client.get_experiment_by_name(experiment_name)
if experiment:
verify_experiment_active(experiment)
experiment_id = experiment.experiment_id
else:
if not experiment:
_logger.info(
"Experiment with name '%s' does not exist. Creating a new experiment.",
experiment_name,
Expand All @@ -123,17 +111,28 @@ def verify_experiment_active(experiment):
# simultaneously, a race condition may be encountered here wherein experiment creation
# fails
experiment_id = client.create_experiment(experiment_name)
Copy link
Collaborator

Choose a reason for hiding this comment

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

There's a potential race condition here between client.get_experiment_by_name() and client.create_experiment(experiment_name). It could cause issues in distributed logging situation. I don't think you need to fix it here, but might be good to annotate the code to make future life easier.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

This is a good call-out. It's worth noting that the MLflow fluent API in general is not meant to be safe across threads or processes. I've added a note here nonetheless.

experiment = client.get_experiment(experiment_id)
else:
experiment = client.get_experiment(experiment_id)
Copy link
Collaborator

Choose a reason for hiding this comment

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

Does this function throw if experiment_id doesn't exist?

Copy link
Member

Choose a reason for hiding this comment

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

404 in REST and throws in file_store and sqlalchemy_store

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Thanks @BenWilson2 ! Yes, while this throws for MLflow-native stores, it may not always throw for alternative stores. I don't think the None check is particularly problematic here as an extra layer of defense.

if experiment is None:
raise MlflowException(
message=f"Experiment with ID '{experiment_id}' does not exist.",
error_code=RESOURCE_DOES_NOT_EXIST,
)
verify_experiment_active(experiment)

if experiment.lifecycle_stage != LifecycleStage.ACTIVE:
raise MlflowException(
message=(
"Cannot set a deleted experiment '%s' as the active experiment."
" You can restore the experiment, or permanently delete the "
" experiment to create a new one." % experiment.name
),
error_code=INVALID_PARAMETER_VALUE,
)

global _active_experiment_id
_active_experiment_id = experiment_id
_active_experiment_id = experiment.experiment_id
return experiment
Copy link
Member

Choose a reason for hiding this comment

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

This actually simplifies some MLOps pipeline flows! A bonus feature is always a great thing



class ActiveRun(Run): # pylint: disable=W0223
Expand Down
11 changes: 6 additions & 5 deletions tests/tracking/test_tracking.py
View file
Open in desktop
Expand Up @@ -80,22 +80,23 @@ def test_create_experiments_with_bad_name_types(name):
def test_set_experiment_by_name():
name = "random_exp"
exp_id = mlflow.create_experiment(name)
mlflow.set_experiment(name)
exp1 = mlflow.set_experiment(name)
assert exp1.experiment_id == exp_id
with start_run() as run:
assert run.info.experiment_id == exp_id

another_name = "another_experiment"
mlflow.set_experiment(another_name)
exp_id2 = mlflow.tracking.MlflowClient().get_experiment_by_name(another_name)
exp2 = mlflow.set_experiment(another_name)
with start_run() as another_run:
assert another_run.info.experiment_id == exp_id2.experiment_id
assert another_run.info.experiment_id == exp2.experiment_id


@pytest.mark.usefixtures("reset_active_experiment")
def test_set_experiment_by_id():
name = "random_exp"
exp_id = mlflow.create_experiment(name)
mlflow.set_experiment(experiment_id=exp_id)
active_exp = mlflow.set_experiment(experiment_id=exp_id)
assert active_exp.experiment_id == exp_id
with start_run() as run:
assert run.info.experiment_id == exp_id

Expand Down