def run_databricks(uri, entry_point, version, parameters, experiment_id,
cluster_spec, git_username, git_password):
"""
Runs the project at the specified URI on Databricks, returning a `SubmittedRun` that can be
used to query the run's status or wait for the resulting Databricks Job run to terminate.
"""
_check_databricks_auth_available()
if cluster_spec is None:
raise ExecutionException(
"Cluster spec must be provided when launching MLflow project runs "
"on Databricks.")
# Fetch the project into work_dir & validate parameters
work_dir = _get_work_dir(uri, use_temp_cwd=True)
_fetch_project(uri, version, work_dir, git_username, git_password)
project = _load_project(work_dir, uri)
project.get_entry_point(entry_point)._validate_parameters(parameters)
# Upload the project to DBFS, get the URI of the project
dbfs_project_uri = _upload_project_to_dbfs(work_dir, experiment_id)
# Create run object with remote tracking server. Get the git commit from the working directory,
# etc.
tracking_uri = tracking.get_tracking_uri()
remote_run = _create_databricks_run(
tracking_uri=tracking_uri,
experiment_id=experiment_id,
source_name=_expand_uri(uri),
source_version=tracking._get_git_commit(work_dir),
entry_point_name=entry_point)
# Set up environment variables for remote execution
env_vars = {}
if experiment_id is not None:
eprint("=== Using experiment ID %s ===" % experiment_id)
env_vars[tracking._EXPERIMENT_ID_ENV_VAR] = experiment_id
if remote_run is not None:
env_vars[tracking._TRACKING_URI_ENV_VAR] = tracking.get_tracking_uri()
env_vars[tracking._RUN_ID_ENV_VAR] = remote_run.run_info.run_uuid
eprint("=== Running entry point %s of project %s on Databricks. ===" %
(entry_point, uri))
# Launch run on Databricks
with open(cluster_spec, 'r') as handle:
try:
cluster_spec = json.load(handle)
except ValueError:
eprint(
"Error when attempting to load and parse JSON cluster spec from file "
"%s. " % cluster_spec)
raise
fuse_dst_dir = os.path.join(
"/dbfs/",
_parse_dbfs_uri_path(dbfs_project_uri).lstrip("/"))
command = _get_databricks_run_cmd(fuse_dst_dir, entry_point, parameters)
db_run_id = _run_shell_command_job(uri, command, env_vars, cluster_spec)
return SubmittedRun(remote_run, DatabricksPollableRun(db_run_id))
def _fetch_and_clean_project(uri, version=None, git_username=None, git_password=None):
"""
Fetches the project at the passed-in URI & prepares it for upload to DBFS. Returns the path of
the temporary directory into which the project was fetched.
"""
work_dir = _fetch_project(
uri=uri, force_tempdir=True, version=version, git_username=git_username,
git_password=git_password)
# Remove the mlruns directory from the fetched project to avoid cache-busting
mlruns_dir = os.path.join(work_dir, "mlruns")
if os.path.exists(mlruns_dir):
shutil.rmtree(mlruns_dir)
return work_dir