Databricks runtime#
The databricks runtime runs on a Databricks cluster (and not in the Iguazio cluster). The function raises a pod on MLRun, which communicates with the Databricks cluster. The requests originate in MLRun and all computing is in the Databricks cluster.
With the databricks runtime, you can send your local file/code as a string to the job, and use a handler as an endpoint for user code. You can optionally send keyword arguments (kwargs) to this job.
You can run the function on:
An existing cluster, by including
DATABRICKS_CLUSTER_ID
A job compute cluster, created and dedicated for this function only.
Params that are not related to a new cluster or an existing cluster:
timeout_minutes
token_key
artifact_json_dir
(location where the json file that contains all logged mlrun artifacts is saved, and which is deleted after the run)
Params that are related to a new cluster:
spark_version
node_type_id
num_workers
Example of a job compute cluster#
To create a job compute cluster, omit DATABRICKS_CLUSTER_ID
, and set the cluster specs by using the task parameters when running the function. For example:
params['task_parameters'] = {'new_cluster_spec': {'node_type_id': 'm5d.large'}, 'number_of_workers': 2, 'timeout_minutes': 15, `token_key`: non-default-value}
Do not send variables named task_parameters
or context
since these are utilized by the internal processes of the runtime.
Example of running a Databricks job from a local file#
This example uses an existing cluster: DATABRICKS_CLUSTER_ID.
import os
import mlrun
from mlrun.runtimes.function_reference import FunctionReference
# If using a Databricks data store, for example, set the credentials:
os.environ["DATABRICKS_HOST"] = "DATABRICKS_HOST"
os.environ["DATABRICKS_TOKEN"] = "DATABRICKS_TOKEN"
os.environ["DATABRICKS_CLUSTER_ID"] = "DATABRICKS_CLUSTER_ID"
def add_databricks_env(function):
job_env = {
"DATABRICKS_HOST": os.environ["DATABRICKS_HOST"],
"DATABRICKS_CLUSTER_ID": os.environ.get("DATABRICKS_CLUSTER_ID"),
}
for name, val in job_env.items():
function.spec.env.append({"name": name, "value": val})
project_name = "databricks-runtime-project"
project = mlrun.get_or_create_project(project_name, context="./", user_project=False)
secrets = {"DATABRICKS_TOKEN": os.environ["DATABRICKS_TOKEN"]}
project.set_secrets(secrets)
code = """
def print_kwargs(**kwargs):
print(f"kwargs: {kwargs}")
"""
function_ref = FunctionReference(
kind="databricks",
code=code,
image="mlrun/mlrun",
name="databricks-function",
)
function = function_ref.to_function()
add_databricks_env(function=function)
run = function.run(
handler="print_kwargs",
project=project_name,
params={
"param1": "value1",
"param2": "value2",
"task_parameters": {"timeout_minutes": 15},
},
)
Logging a Databricks response as an artifact#
import numpy as np
import pandas as pd
from pyspark.sql import SparkSession
def main():
df = pd.DataFrame({"A": np.random.randint(1, 100, 5), "B": np.random.rand(5)})
path = "/dbfs/path/folder"
parquet_df_path = f"{path}/df.parquet"
csv_df_path = f"{path}/df.csv"
if not os.path.exists(path):
os.makedirs(path)
# save df
df.to_parquet(parquet_df_path)
df.to_csv(csv_df_path, index=False)
# log artifact
mlrun_log_artifact("parquet_artifact", parquet_df_path)
mlrun_log_artifact("csv_artifact", csv_df_path)
# spark
spark = SparkSession.builder.appName("example").getOrCreate()
spark_df = spark.createDataFrame(df)
# spark path format:
spark_parquet_path = "dbfs:///path/folder/spark_df.parquet"
spark_df.write.mode("overwrite").parquet(spark_parquet_path)
mlrun_log_artifact("spark_artifact", spark_parquet_path)
# an illegal artifact does not raise an error, it logs an error log instead, for example:
# mlrun_log_artifact("illegal_artifact", "/not_exists_path/illegal_df.parquet")
function = mlrun.code_to_function(
name="databricks-log_artifact",
kind="databricks",
project=project_name,
filename="./databricks_job.py",
image="mlrun/mlrun",
)
add_databricks_env(function=function)
run = function.run(
handler="main",
project=project_name,
)
project.list_artifacts()