chore: import upstream snapshot with attribution

examples/databricks/dbconnect.py

@@ -0,0 +1,56 @@
+"""
+python examples/databricks/dbconnect.py --cluster-id <cluster-id>
+"""
+
+import argparse
+
+from databricks.connect import DatabricksSession
+from databricks.sdk import WorkspaceClient
+from pyspark.sql.types import DoubleType
+from sklearn import datasets
+from sklearn.neighbors import KNeighborsClassifier
+
+import mlflow
+from mlflow.models import infer_signature
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--cluster-id", required=True)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    wc = WorkspaceClient()
+
+    # Train a model
+    X, y = datasets.load_iris(as_frame=True, return_X_y=True)
+    model = KNeighborsClassifier().fit(X, y)
+    predictions = model.predict(X)
+    signature = infer_signature(X, predictions)
+
+    # Log the model
+    mlflow.set_tracking_uri("databricks")
+    mlflow.set_experiment(f"/Users/{wc.current_user.me().user_name}/dbconnect")
+    with mlflow.start_run():
+        model_info = mlflow.sklearn.log_model(model, name="model", signature=signature)
+
+    spark = DatabricksSession.builder.remote(
+        host=wc.config.host,
+        token=wc.config.token,
+        cluster_id=args.cluster_id,
+    ).getOrCreate()
+    sdf = spark.createDataFrame(X.head(5))
+    pyfunc_udf = mlflow.pyfunc.spark_udf(
+        spark,
+        model_info.model_uri,
+        env_manager="local",
+        result_type=DoubleType(),
+    )
+    preds = sdf.select(pyfunc_udf(*X.columns).alias("preds"))
+    preds.show()
+
+
+if __name__ == "__main__":
+    main()

examples/databricks/log_runs.py

@@ -0,0 +1,56 @@
+"""
+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()

examples/databricks/multipart.py

@@ -0,0 +1,148 @@
+"""
+Benchmark for multi-part upload and download of artifacts.
+"""
+
+import hashlib
+import json
+import os
+import pathlib
+import tempfile
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+import pandas as pd
+import psutil
+from tqdm.auto import tqdm
+
+import mlflow
+from mlflow.environment_variables import (
+    MLFLOW_ENABLE_MULTIPART_DOWNLOAD,
+    MLFLOW_ENABLE_MULTIPART_UPLOAD,
+)
+from mlflow.utils.time import Timer
+
+GiB = 1024**3
+
+
+def show_system_info():
+    svmem = psutil.virtual_memory()
+    info = json.dumps(
+        {
+            "MLflow version": mlflow.__version__,
+            "MPU enabled": MLFLOW_ENABLE_MULTIPART_DOWNLOAD.get(),
+            "MPD enabled": MLFLOW_ENABLE_MULTIPART_UPLOAD.get(),
+            "CPU count": psutil.cpu_count(),
+            "Memory usage (total) [GiB]": svmem.total // GiB,
+            "Memory used [GiB]": svmem.used // GiB,
+            "Memory available [GiB]": svmem.available // GiB,
+        },
+        indent=2,
+    )
+    max_len = max(map(len, info.splitlines()))
+    print("=" * max_len)
+    print(info)
+    print("=" * max_len)
+
+
+def md5_checksum(path):
+    file_hash = hashlib.sha256()
+    with open(path, "rb") as f:
+        while chunk := f.read(1024**2):
+            file_hash.update(chunk)
+    return file_hash.hexdigest()
+
+
+def assert_checksum_equal(path1, path2):
+    assert md5_checksum(path1) == md5_checksum(path2), f"Checksum mismatch for {path1} and {path2}"
+
+
+def yield_random_bytes(num_bytes):
+    while num_bytes > 0:
+        chunk_size = min(num_bytes, 1024**2)
+        yield os.urandom(chunk_size)
+        num_bytes -= chunk_size
+
+
+def generate_random_file(path, num_bytes):
+    with open(path, "wb") as f:
+        for chunk in yield_random_bytes(num_bytes):
+            f.write(chunk)
+
+
+def upload_and_download(file_size, num_files):
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tmpdir = pathlib.Path(tmpdir)
+
+        # Prepare files
+        src_dir = tmpdir / "src"
+        src_dir.mkdir()
+        files = {}
+        with ThreadPoolExecutor() as pool:
+            futures = []
+            for i in range(num_files):
+                f = src_dir / str(i)
+                futures.append(pool.submit(generate_random_file, f, file_size))
+                files[f.name] = f
+
+            for fut in tqdm(
+                as_completed(futures),
+                total=len(futures),
+                desc="Generating files",
+                colour="#FFA500",
+            ):
+                fut.result()
+
+        # Upload
+        with mlflow.start_run() as run:
+            with Timer() as t_upload:
+                mlflow.log_artifacts(str(src_dir))
+
+        # Download
+        dst_dir = tmpdir / "dst"
+        dst_dir.mkdir()
+        with Timer() as t_download:
+            mlflow.artifacts.download_artifacts(
+                artifact_uri=f"{run.info.artifact_uri}/", dst_path=dst_dir
+            )
+
+        # Verify checksums
+        with ThreadPoolExecutor() as pool:
+            futures = []
+            for f in dst_dir.rglob("*"):
+                if f.is_dir():
+                    continue
+                futures.append(pool.submit(assert_checksum_equal, f, files[f.name]))
+
+            for fut in tqdm(
+                as_completed(futures),
+                total=len(futures),
+                desc="Verifying checksums",
+                colour="#FFA500",
+            ):
+                fut.result()
+
+        return t_upload.elapsed, t_download.elapsed
+
+
+def main():
+    # Uncomment the following lines if you're running this script outside of Databricks
+    # using a personal access token:
+    # mlflow.set_tracking_uri("databricks")
+    # mlflow.set_experiment("/Users/<username>/benchmark")
+
+    FILE_SIZE = 1 * GiB
+    NUM_FILES = 2
+    NUM_ATTEMPTS = 3
+
+    show_system_info()
+    stats = []
+    for i in range(NUM_ATTEMPTS):
+        print(f"Attempt {i + 1} / {NUM_ATTEMPTS}")
+        stats.append(upload_and_download(FILE_SIZE, NUM_FILES))
+
+    df = pd.DataFrame(stats, columns=["upload [s]", "download [s]"])
+    # show mean, min, max in markdown table
+    print(df.aggregate(["count", "mean", "min", "max"]).to_markdown())
+
+
+if __name__ == "__main__":
+    main()