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.
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
Add optional experiment_id parameter to mlflow.set_experiment
#5012
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -14,6 +14,10 @@ | |
from mlflow.entities import Experiment, Run, RunInfo, RunStatus, Param, RunTag, Metric, ViewType | ||
from mlflow.entities.lifecycle_stage import LifecycleStage | ||
from mlflow.exceptions import MlflowException | ||
from mlflow.protos.databricks_pb2 import ( | ||
INVALID_PARAMETER_VALUE, | ||
RESOURCE_DOES_NOT_EXIST, | ||
) | ||
from mlflow.tracking.client import MlflowClient | ||
from mlflow.tracking import artifact_utils, _get_store | ||
from mlflow.tracking.context import registry as context_registry | ||
|
@@ -50,12 +54,19 @@ | |
_logger = logging.getLogger(__name__) | ||
|
||
|
||
def set_experiment(experiment_name: str) -> None: | ||
def set_experiment(experiment_name: str = None, experiment_id: str = None) -> None: | ||
""" | ||
Set given experiment as active experiment. If experiment does not exist, create an experiment | ||
with provided name. | ||
|
||
:param experiment_name: Case sensitive name of an experiment to be activated. | ||
Set the given experiment as the active experiment. The experiment must either be specified by | ||
name via `experiment_name` or by ID via `experiment_id`. The experiment name and ID cannot | ||
both be specified. | ||
|
||
:param experiment_name: Case sensitive name of the experiment to be activated. If an experiment | ||
with this name does not exist, a new experiment wth this name is | ||
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 | ||
|
@@ -80,20 +91,48 @@ def set_experiment(experiment_name: str) -> None: | |
Tags: {} | ||
Lifecycle_stage: active | ||
""" | ||
if (experiment_name is not None and experiment_id is not None) or ( | ||
experiment_name is None and experiment_id is None | ||
): | ||
raise MlflowException( | ||
message="Must specify exactly one of: `experiment_id` or `experiment_name`.", | ||
error_code=INVALID_PARAMETER_VALUE, | ||
) | ||
|
||
client = MlflowClient() | ||
experiment = client.get_experiment_by_name(experiment_name) | ||
exp_id = experiment.experiment_id if experiment else None | ||
if exp_id is None: # id can be 0 | ||
print("INFO: '{}' does not exist. Creating a new experiment".format(experiment_name)) | ||
exp_id = client.create_experiment(experiment_name) | ||
elif experiment.lifecycle_stage == LifecycleStage.DELETED: | ||
if experiment_id is None: | ||
experiment = client.get_experiment_by_name(experiment_name) | ||
if not experiment: | ||
_logger.info( | ||
"Experiment with name '%s' does not exist. Creating a new experiment.", | ||
experiment_name, | ||
) | ||
# NB: If two simultaneous threads or processes attempt to set the same experiment | ||
# simultaneously, a race condition may be encountered here wherein experiment creation | ||
# fails | ||
experiment_id = client.create_experiment(experiment_name) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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. There was a problem hiding this comment. Choose a reason for hiding this commentThe 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) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Does this function throw if experiment_id doesn't exist? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 404 in REST and throws in file_store and sqlalchemy_store There was a problem hiding this comment. Choose a reason for hiding this commentThe 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 |
||
if experiment is None: | ||
raise MlflowException( | ||
message=f"Experiment with ID '{experiment_id}' does not exist.", | ||
error_code=RESOURCE_DOES_NOT_EXIST, | ||
) | ||
|
||
if experiment.lifecycle_stage != LifecycleStage.ACTIVE: | ||
raise MlflowException( | ||
"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 | ||
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 = exp_id | ||
_active_experiment_id = experiment.experiment_id | ||
return experiment | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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 | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Is empty string legit for either name or id? I guess not. In that case, it might be better to verify
if not (experiment_name and experiment_id)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I am curious on why do we need this condition? We just need
experiment_name is None and experiment_id is None
right? Let me know if I am missing something.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Wouldn't the guarded exclusivity check protect against conflicting submission behavior. For an invalid id but a valid name it would raise an Exception based on the id validation. Could be super confusing for users.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I would imagine that it's invalid on most if not all backends, but I'm hesitant to enforce this on the off chance that someone's third-party backend has a legitimate use case for this. For example, MLflow's default experiment used to be a falsey integer value (
0
) before the change to string representations.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Because the same logic applies for experiment names, I've removed the test case asserting that an empty string name is invalid, as that's backend-dependent.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
SGTM.