wehub-resource-sync
71 lines
2.0 KiB
Python
71 lines
2.0 KiB
Python
_ERROR_MSG = (
|
|
"The `databricks-agents` package is required to use databricks trace archival. "
|
|
"Please install it with `pip install databricks-agents`."
|
|
)
|
|
|
|
|
|
def enable_databricks_trace_archival(
|
|
*,
|
|
delta_table_fullname: str,
|
|
experiment_id: str | None = None,
|
|
) -> None:
|
|
"""
|
|
Enable archiving traces for an MLflow experiment to a Unity Catalog Delta table.
|
|
|
|
Args:
|
|
delta_table_fullname: The full name of the Unity Catalog Delta table to archive traces to.
|
|
experiment_id: The MLflow experiment ID to enable archival for.
|
|
Default to the current active experiment.
|
|
|
|
Example:
|
|
|
|
.. code-block:: python
|
|
|
|
from mlflow.tracing.archival import enable_databricks_trace_archival
|
|
|
|
enable_databricks_trace_archival(
|
|
delta_table_fullname="my_catalog.my_schema.my_prefix",
|
|
experiment_id="12345",
|
|
)
|
|
"""
|
|
from mlflow.tracking.fluent import _get_experiment_id
|
|
|
|
try:
|
|
from databricks.agents.archive import enable_trace_archival
|
|
except ImportError:
|
|
raise ImportError(_ERROR_MSG)
|
|
|
|
experiment_id = experiment_id or _get_experiment_id()
|
|
|
|
enable_trace_archival(
|
|
experiment_id=experiment_id,
|
|
table_fullname=delta_table_fullname,
|
|
)
|
|
|
|
|
|
def disable_databricks_trace_archival(*, experiment_id: str | None = None) -> None:
|
|
"""
|
|
Disable archiving traces for an MLflow experiment to a Unity Catalog Delta table.
|
|
|
|
Args:
|
|
experiment_id: The MLflow experiment ID to disable archival for.
|
|
|
|
Example:
|
|
|
|
.. code-block:: python
|
|
|
|
from mlflow.tracing.archival import disable_databricks_trace_archival
|
|
|
|
disable_databricks_trace_archival(experiment_id="12345")
|
|
"""
|
|
from mlflow.tracking.fluent import _get_experiment_id
|
|
|
|
try:
|
|
from databricks.agents.archive import disable_trace_archival
|
|
except ImportError:
|
|
raise ImportError(_ERROR_MSG)
|
|
|
|
experiment_id = experiment_id or _get_experiment_id()
|
|
|
|
disable_trace_archival(experiment_id=experiment_id)
|