wehub-resource-sync
57 lines
1.7 KiB
Python
57 lines
1.7 KiB
Python
"""
|
|
Logs MLflow runs in Databricks from an external host.
|
|
|
|
How to run:
|
|
$ python examples/databricks/log_runs.py --host <host> --token <token> --user <user> [--experiment-id 123]
|
|
|
|
See also:
|
|
https://docs.databricks.com/dev-tools/api/latest/authentication.html#generate-a-personal-access-token
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import uuid
|
|
|
|
from sklearn import datasets, svm
|
|
from sklearn.model_selection import GridSearchCV, ParameterGrid
|
|
|
|
import mlflow
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--host", help="Databricks workspace URL")
|
|
parser.add_argument("--token", help="Databricks personal access token")
|
|
parser.add_argument("--user", help="Databricks username")
|
|
parser.add_argument(
|
|
"--experiment-id",
|
|
default=None,
|
|
help="ID of the experiment to log runs in. If unspecified, a new experiment will be created.",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
os.environ["DATABRICKS_HOST"] = args.host
|
|
os.environ["DATABRICKS_TOKEN"] = args.token
|
|
|
|
mlflow.set_tracking_uri("databricks")
|
|
if args.experiment_id:
|
|
experiment = mlflow.set_experiment(experiment_id=args.experiment_id)
|
|
else:
|
|
experiment = mlflow.set_experiment(f"/Users/{args.user}/{uuid.uuid4().hex}")
|
|
|
|
print(f"Logging runs in {args.host}#/mlflow/experiments/{experiment.experiment_id}")
|
|
mlflow.sklearn.autolog(max_tuning_runs=None)
|
|
iris = datasets.load_iris()
|
|
parameters = {"kernel": ("linear", "rbf"), "C": [1, 5, 10]}
|
|
clf = GridSearchCV(svm.SVC(), parameters)
|
|
clf.fit(iris.data, iris.target)
|
|
|
|
# Log unnested runs
|
|
for params in ParameterGrid(parameters):
|
|
clf = svm.SVC(**params)
|
|
clf.fit(iris.data, iris.target)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|