347 lines
9.9 KiB
Python
347 lines
9.9 KiB
Python
"""
|
|
The ``mlflow`` module provides a high-level "fluent" API for starting and managing MLflow runs.
|
|
For example:
|
|
|
|
.. code:: python
|
|
|
|
import mlflow
|
|
|
|
mlflow.start_run()
|
|
mlflow.log_param("my", "param")
|
|
mlflow.log_metric("score", 100)
|
|
mlflow.end_run()
|
|
|
|
You can also use the context manager syntax like this:
|
|
|
|
.. code:: python
|
|
|
|
with mlflow.start_run() as run:
|
|
mlflow.log_param("my", "param")
|
|
mlflow.log_metric("score", 100)
|
|
|
|
which automatically terminates the run at the end of the ``with`` block.
|
|
|
|
The fluent tracking API is not currently threadsafe. Any concurrent callers to the tracking API must
|
|
implement mutual exclusion manually.
|
|
|
|
For a lower level API, see the :py:mod:`mlflow.client` module.
|
|
"""
|
|
|
|
import contextlib
|
|
from typing import TYPE_CHECKING
|
|
|
|
from mlflow.version import VERSION
|
|
|
|
__version__ = VERSION
|
|
|
|
import mlflow.mismatch
|
|
|
|
# `check_version_mismatch` must be called here before importing any other modules
|
|
with contextlib.suppress(Exception):
|
|
mlflow.mismatch._check_version_mismatch()
|
|
|
|
from mlflow import (
|
|
artifacts, # noqa: F401
|
|
client, # noqa: F401
|
|
config, # noqa: F401
|
|
data, # noqa: F401
|
|
exceptions, # noqa: F401
|
|
models, # noqa: F401
|
|
projects, # noqa: F401
|
|
tracing, # noqa: F401
|
|
tracking, # noqa: F401
|
|
)
|
|
from mlflow.environment_variables import MLFLOW_CONFIGURE_LOGGING
|
|
from mlflow.utils.lazy_load import LazyLoader
|
|
from mlflow.utils.logging_utils import _configure_mlflow_loggers
|
|
|
|
# Lazily load mlflow flavors to avoid excessive dependencies.
|
|
anthropic = LazyLoader("mlflow.anthropic", globals(), "mlflow.anthropic")
|
|
autogen = LazyLoader("mlflow.autogen", globals(), "mlflow.autogen")
|
|
bedrock = LazyLoader("mlflow.bedrock", globals(), "mlflow.bedrock")
|
|
catboost = LazyLoader("mlflow.catboost", globals(), "mlflow.catboost")
|
|
crewai = LazyLoader("mlflow.crewai", globals(), "mlflow.crewai")
|
|
diviner = LazyLoader("mlflow.diviner", globals(), "mlflow.diviner")
|
|
dspy = LazyLoader("mlflow.dspy", globals(), "mlflow.dspy")
|
|
fastai = LazyLoader("mlflow.fastai", globals(), "mlflow.fastai")
|
|
gemini = LazyLoader("mlflow.gemini", globals(), "mlflow.gemini")
|
|
groq = LazyLoader("mlflow.groq", globals(), "mlflow.groq")
|
|
h2o = LazyLoader("mlflow.h2o", globals(), "mlflow.h2o")
|
|
johnsnowlabs = LazyLoader("mlflow.johnsnowlabs", globals(), "mlflow.johnsnowlabs")
|
|
keras = LazyLoader("mlflow.keras", globals(), "mlflow.keras")
|
|
langchain = LazyLoader("mlflow.langchain", globals(), "mlflow.langchain")
|
|
lightgbm = LazyLoader("mlflow.lightgbm", globals(), "mlflow.lightgbm")
|
|
litellm = LazyLoader("mlflow.litellm", globals(), "mlflow.litellm")
|
|
llama_index = LazyLoader("mlflow.llama_index", globals(), "mlflow.llama_index")
|
|
llm = LazyLoader("mlflow.llm", globals(), "mlflow.llm")
|
|
metrics = LazyLoader("mlflow.metrics", globals(), "mlflow.metrics")
|
|
mistral = LazyLoader("mlflow.mistral", globals(), "mlflow.mistral")
|
|
mleap = LazyLoader("mlflow.mleap", globals(), "mlflow.mleap")
|
|
onnx = LazyLoader("mlflow.onnx", globals(), "mlflow.onnx")
|
|
openai = LazyLoader("mlflow.openai", globals(), "mlflow.openai")
|
|
paddle = LazyLoader("mlflow.paddle", globals(), "mlflow.paddle")
|
|
pmdarima = LazyLoader("mlflow.pmdarima", globals(), "mlflow.pmdarima")
|
|
promptflow = LazyLoader("mlflow.promptflow", globals(), "mlflow.promptflow")
|
|
prophet = LazyLoader("mlflow.prophet", globals(), "mlflow.prophet")
|
|
pyfunc = LazyLoader("mlflow.pyfunc", globals(), "mlflow.pyfunc")
|
|
pyspark = LazyLoader("mlflow.pyspark", globals(), "mlflow.pyspark")
|
|
pytorch = LazyLoader("mlflow.pytorch", globals(), "mlflow.pytorch")
|
|
rfunc = LazyLoader("mlflow.rfunc", globals(), "mlflow.rfunc")
|
|
recipes = LazyLoader("mlflow.recipes", globals(), "mlflow.recipes")
|
|
sentence_transformers = LazyLoader(
|
|
"mlflow.sentence_transformers",
|
|
globals(),
|
|
"mlflow.sentence_transformers",
|
|
)
|
|
shap = LazyLoader("mlflow.shap", globals(), "mlflow.shap")
|
|
sklearn = LazyLoader("mlflow.sklearn", globals(), "mlflow.sklearn")
|
|
spacy = LazyLoader("mlflow.spacy", globals(), "mlflow.spacy")
|
|
spark = LazyLoader("mlflow.spark", globals(), "mlflow.spark")
|
|
statsmodels = LazyLoader("mlflow.statsmodels", globals(), "mlflow.statsmodels")
|
|
tensorflow = LazyLoader("mlflow.tensorflow", globals(), "mlflow.tensorflow")
|
|
# TxtAI integration is defined at https://github.com/neuml/mlflow-txtai
|
|
txtai = LazyLoader("mlflow.txtai", globals(), "mlflow_txtai")
|
|
transformers = LazyLoader("mlflow.transformers", globals(), "mlflow.transformers")
|
|
xgboost = LazyLoader("mlflow.xgboost", globals(), "mlflow.xgboost")
|
|
|
|
if TYPE_CHECKING:
|
|
# Do not move this block above the lazy-loaded modules above.
|
|
# All the lazy-loaded modules above must be imported here for code completion to work in IDEs.
|
|
from mlflow import ( # noqa: F401
|
|
anthropic,
|
|
autogen,
|
|
bedrock,
|
|
catboost,
|
|
crewai,
|
|
diviner,
|
|
dspy,
|
|
fastai,
|
|
gemini,
|
|
groq,
|
|
h2o,
|
|
johnsnowlabs,
|
|
keras,
|
|
langchain,
|
|
lightgbm,
|
|
litellm,
|
|
llama_index,
|
|
llm,
|
|
metrics,
|
|
mistral,
|
|
mleap,
|
|
onnx,
|
|
openai,
|
|
paddle,
|
|
pmdarima,
|
|
promptflow,
|
|
prophet,
|
|
pyfunc,
|
|
pyspark,
|
|
pytorch,
|
|
recipes,
|
|
rfunc,
|
|
sentence_transformers,
|
|
shap,
|
|
sklearn,
|
|
spacy,
|
|
spark,
|
|
statsmodels,
|
|
tensorflow,
|
|
transformers,
|
|
xgboost,
|
|
)
|
|
|
|
if MLFLOW_CONFIGURE_LOGGING.get() is True:
|
|
_configure_mlflow_loggers(root_module_name=__name__)
|
|
|
|
from mlflow.client import MlflowClient
|
|
|
|
# For backward compatibility, we expose the following functions and classes at the top level in
|
|
# addition to `mlflow.config`.
|
|
from mlflow.config import (
|
|
disable_system_metrics_logging,
|
|
enable_system_metrics_logging,
|
|
get_registry_uri,
|
|
get_tracking_uri,
|
|
is_tracking_uri_set,
|
|
set_registry_uri,
|
|
set_system_metrics_node_id,
|
|
set_system_metrics_samples_before_logging,
|
|
set_system_metrics_sampling_interval,
|
|
set_tracking_uri,
|
|
)
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.models import evaluate
|
|
from mlflow.models.evaluation.validation import validate_evaluation_results
|
|
from mlflow.projects import run
|
|
from mlflow.tracing.assessment import (
|
|
delete_expectation,
|
|
delete_feedback,
|
|
log_expectation,
|
|
log_feedback,
|
|
update_expectation,
|
|
update_feedback,
|
|
)
|
|
from mlflow.tracing.fluent import (
|
|
add_trace,
|
|
get_current_active_span,
|
|
get_last_active_trace,
|
|
get_last_active_trace_id,
|
|
get_trace,
|
|
log_trace,
|
|
search_traces,
|
|
start_span,
|
|
trace,
|
|
update_current_trace,
|
|
)
|
|
from mlflow.tracking._model_registry.fluent import (
|
|
delete_prompt,
|
|
delete_prompt_alias,
|
|
load_prompt,
|
|
register_model,
|
|
register_prompt,
|
|
search_model_versions,
|
|
search_registered_models,
|
|
set_prompt_alias,
|
|
)
|
|
from mlflow.tracking.fluent import (
|
|
ActiveRun,
|
|
active_run,
|
|
autolog,
|
|
create_experiment,
|
|
delete_experiment,
|
|
delete_run,
|
|
delete_tag,
|
|
end_run,
|
|
flush_artifact_async_logging,
|
|
flush_async_logging,
|
|
flush_trace_async_logging,
|
|
get_artifact_uri,
|
|
get_experiment,
|
|
get_experiment_by_name,
|
|
get_parent_run,
|
|
get_run,
|
|
last_active_run,
|
|
load_table,
|
|
log_artifact,
|
|
log_artifacts,
|
|
log_dict,
|
|
log_figure,
|
|
log_image,
|
|
log_input,
|
|
log_metric,
|
|
log_metrics,
|
|
log_param,
|
|
log_params,
|
|
log_table,
|
|
log_text,
|
|
search_experiments,
|
|
search_runs,
|
|
set_experiment,
|
|
set_experiment_tag,
|
|
set_experiment_tags,
|
|
set_tag,
|
|
set_tags,
|
|
start_run,
|
|
)
|
|
from mlflow.tracking.multimedia import Image
|
|
from mlflow.utils.async_logging.run_operations import RunOperations # noqa: F401
|
|
from mlflow.utils.credentials import login
|
|
from mlflow.utils.doctor import doctor
|
|
|
|
__all__ = [
|
|
"ActiveRun",
|
|
"MlflowClient",
|
|
"MlflowException",
|
|
"active_run",
|
|
"autolog",
|
|
"create_experiment",
|
|
"delete_experiment",
|
|
"delete_run",
|
|
"delete_tag",
|
|
"disable_system_metrics_logging",
|
|
"doctor",
|
|
"enable_system_metrics_logging",
|
|
"end_run",
|
|
"evaluate",
|
|
"flush_async_logging",
|
|
"flush_artifact_async_logging",
|
|
"flush_trace_async_logging",
|
|
"get_artifact_uri",
|
|
"get_experiment",
|
|
"get_experiment_by_name",
|
|
"get_last_active_trace",
|
|
"get_last_active_trace_id",
|
|
"get_parent_run",
|
|
"get_registry_uri",
|
|
"get_run",
|
|
"get_tracking_uri",
|
|
"is_tracking_uri_set",
|
|
"last_active_run",
|
|
"load_table",
|
|
"log_artifact",
|
|
"log_artifacts",
|
|
"log_dict",
|
|
"log_figure",
|
|
"log_image",
|
|
"log_input",
|
|
"log_metric",
|
|
"log_metrics",
|
|
"log_param",
|
|
"log_params",
|
|
"log_table",
|
|
"log_text",
|
|
"login",
|
|
"pyfunc",
|
|
"register_model",
|
|
"run",
|
|
"search_experiments",
|
|
"search_model_versions",
|
|
"search_registered_models",
|
|
"search_runs",
|
|
"set_experiment",
|
|
"set_experiment_tag",
|
|
"set_experiment_tags",
|
|
"set_registry_uri",
|
|
"set_system_metrics_node_id",
|
|
"set_system_metrics_samples_before_logging",
|
|
"set_system_metrics_sampling_interval",
|
|
"set_tag",
|
|
"set_tags",
|
|
"set_tracking_uri",
|
|
"start_run",
|
|
"validate_evaluation_results",
|
|
"Image",
|
|
# Tracing Fluent APIs
|
|
"get_current_active_span",
|
|
"get_trace",
|
|
"search_traces",
|
|
"start_span",
|
|
"trace",
|
|
"add_trace",
|
|
"log_trace",
|
|
"update_current_trace",
|
|
# Assessment APIs
|
|
"delete_expectation",
|
|
"delete_feedback",
|
|
"log_expectation",
|
|
"log_feedback",
|
|
"update_expectation",
|
|
"update_feedback",
|
|
# Prompt Registry APIs
|
|
"delete_prompt",
|
|
"load_prompt",
|
|
"register_prompt",
|
|
"set_prompt_alias",
|
|
"delete_prompt_alias",
|
|
]
|
|
|
|
|
|
# `mlflow.gateway` depends on optional dependencies such as pydantic, psutil, and has version
|
|
# restrictions for dependencies. Importing this module fails if they are not installed or
|
|
# if invalid versions of these required packages are installed.
|
|
with contextlib.suppress(Exception):
|
|
from mlflow import gateway # noqa: F401
|
|
|
|
__all__.append("gateway")
|