-
Notifications
You must be signed in to change notification settings - Fork 4.1k
/
local.py
424 lines (380 loc) · 16.7 KB
/
local.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
import logging
import os
import platform
import posixpath
import subprocess
import sys
from pathlib import Path
import mlflow
from mlflow import tracking
from mlflow.environment_variables import (
MLFLOW_KERBEROS_TICKET_CACHE,
MLFLOW_KERBEROS_USER,
MLFLOW_PYARROW_EXTRA_CONF,
)
from mlflow.exceptions import MlflowException
from mlflow.projects import env_type
from mlflow.projects.backend.abstract_backend import AbstractBackend
from mlflow.projects.submitted_run import LocalSubmittedRun
from mlflow.projects.utils import (
MLFLOW_DOCKER_WORKDIR_PATH,
MLFLOW_LOCAL_BACKEND_RUN_ID_CONFIG,
PROJECT_BUILD_IMAGE,
PROJECT_DOCKER_ARGS,
PROJECT_DOCKER_AUTH,
PROJECT_ENV_MANAGER,
PROJECT_STORAGE_DIR,
PROJECT_SYNCHRONOUS,
fetch_and_validate_project,
get_entry_point_command,
get_or_create_run,
get_run_env_vars,
load_project,
)
from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository
from mlflow.store.artifact.azure_blob_artifact_repo import AzureBlobArtifactRepository
from mlflow.store.artifact.gcs_artifact_repo import GCSArtifactRepository
from mlflow.store.artifact.hdfs_artifact_repo import HdfsArtifactRepository
from mlflow.store.artifact.local_artifact_repo import LocalArtifactRepository
from mlflow.store.artifact.s3_artifact_repo import S3ArtifactRepository
from mlflow.utils import env_manager as _EnvManager
from mlflow.utils.conda import get_or_create_conda_env
from mlflow.utils.databricks_utils import get_databricks_env_vars, is_in_databricks_runtime
from mlflow.utils.environment import _PythonEnv
from mlflow.utils.file_utils import get_or_create_nfs_tmp_dir
from mlflow.utils.mlflow_tags import MLFLOW_PROJECT_ENV
from mlflow.utils.os import is_windows
from mlflow.utils.virtualenv import (
_PYENV_ROOT_DIR,
_VIRTUALENV_ENVS_DIR,
_create_virtualenv,
_get_mlflow_virtualenv_root,
_get_virtualenv_extra_env_vars,
_get_virtualenv_name,
_install_python,
)
_logger = logging.getLogger(__name__)
def _env_type_to_env_manager(env_typ):
if env_typ == env_type.CONDA:
return _EnvManager.CONDA
elif env_typ == env_type.PYTHON:
return _EnvManager.VIRTUALENV
elif env_typ == env_type.DOCKER:
return _EnvManager.LOCAL
class LocalBackend(AbstractBackend):
def run(
self,
project_uri,
entry_point,
params,
version,
backend_config,
tracking_uri,
experiment_id,
):
work_dir = fetch_and_validate_project(project_uri, version, entry_point, params)
project = load_project(work_dir)
if MLFLOW_LOCAL_BACKEND_RUN_ID_CONFIG in backend_config:
run_id = backend_config[MLFLOW_LOCAL_BACKEND_RUN_ID_CONFIG]
else:
run_id = None
active_run = get_or_create_run(
run_id, project_uri, experiment_id, work_dir, version, entry_point, params
)
command_args = []
command_separator = " "
env_manager = backend_config[PROJECT_ENV_MANAGER]
synchronous = backend_config[PROJECT_SYNCHRONOUS]
docker_args = backend_config[PROJECT_DOCKER_ARGS]
storage_dir = backend_config[PROJECT_STORAGE_DIR]
build_image = backend_config[PROJECT_BUILD_IMAGE]
docker_auth = backend_config[PROJECT_DOCKER_AUTH]
# Select an appropriate env manager for the project env type
if env_manager is None:
env_manager = _env_type_to_env_manager(project.env_type)
else:
if project.env_type == env_type.PYTHON and env_manager == _EnvManager.CONDA:
raise MlflowException.invalid_parameter_value(
"python_env project cannot be executed using conda. Set `--env-manager` to "
"'virtualenv' or 'local' to execute this project."
)
# If a docker_env attribute is defined in MLproject then it takes precedence over conda yaml
# environments, so the project will be executed inside a docker container.
if project.docker_env:
from mlflow.projects.docker import (
build_docker_image,
validate_docker_env,
validate_docker_installation,
)
tracking.MlflowClient().set_tag(active_run.info.run_id, MLFLOW_PROJECT_ENV, "docker")
validate_docker_env(project)
validate_docker_installation()
image = build_docker_image(
work_dir=work_dir,
repository_uri=project.name,
base_image=project.docker_env.get("image"),
run_id=active_run.info.run_id,
build_image=build_image,
docker_auth=docker_auth,
)
command_args += _get_docker_command(
image=image,
active_run=active_run,
docker_args=docker_args,
volumes=project.docker_env.get("volumes"),
user_env_vars=project.docker_env.get("environment"),
)
# Synchronously create a conda environment (even though this may take some time)
# to avoid failures due to multiple concurrent attempts to create the same conda env.
elif env_manager == _EnvManager.VIRTUALENV:
tracking.MlflowClient().set_tag(
active_run.info.run_id, MLFLOW_PROJECT_ENV, "virtualenv"
)
command_separator = " && "
if project.env_type == env_type.CONDA:
python_env = _PythonEnv.from_conda_yaml(project.env_config_path)
else:
python_env = (
_PythonEnv.from_yaml(project.env_config_path)
if project.env_config_path
else _PythonEnv()
)
if is_in_databricks_runtime():
nfs_tmp_dir = get_or_create_nfs_tmp_dir()
env_root = Path(nfs_tmp_dir) / "envs"
pyenv_root = env_root / _PYENV_ROOT_DIR
virtualenv_root = env_root / _VIRTUALENV_ENVS_DIR
env_vars = _get_virtualenv_extra_env_vars(str(env_root))
else:
pyenv_root = None
virtualenv_root = Path(_get_mlflow_virtualenv_root())
env_vars = None
python_bin_path = _install_python(python_env.python, pyenv_root=pyenv_root)
work_dir_path = Path(work_dir)
env_name = _get_virtualenv_name(python_env, work_dir_path)
env_dir = virtualenv_root / env_name
activate_cmd = _create_virtualenv(
work_dir_path, python_bin_path, env_dir, python_env, extra_env=env_vars
)
command_args += [activate_cmd]
elif env_manager == _EnvManager.CONDA:
tracking.MlflowClient().set_tag(active_run.info.run_id, MLFLOW_PROJECT_ENV, "conda")
command_separator = " && "
conda_env = get_or_create_conda_env(project.env_config_path)
command_args += conda_env.get_activate_command()
# In synchronous mode, run the entry point command in a blocking fashion, sending status
# updates to the tracking server when finished. Note that the run state may not be
# persisted to the tracking server if interrupted
if synchronous:
command_args += get_entry_point_command(project, entry_point, params, storage_dir)
command_str = command_separator.join(command_args)
return _run_entry_point(
command_str, work_dir, experiment_id, run_id=active_run.info.run_id
)
# Otherwise, invoke `mlflow run` in a subprocess
return _invoke_mlflow_run_subprocess(
work_dir=work_dir,
entry_point=entry_point,
parameters=params,
experiment_id=experiment_id,
env_manager=env_manager,
docker_args=docker_args,
storage_dir=storage_dir,
run_id=active_run.info.run_id,
)
def _invoke_mlflow_run_subprocess(
work_dir, entry_point, parameters, experiment_id, env_manager, docker_args, storage_dir, run_id
):
"""
Run an MLflow project asynchronously by invoking ``mlflow run`` in a subprocess, returning
a SubmittedRun that can be used to query run status.
"""
_logger.info("=== Asynchronously launching MLflow run with ID %s ===", run_id)
mlflow_run_arr = _build_mlflow_run_cmd(
uri=work_dir,
entry_point=entry_point,
docker_args=docker_args,
storage_dir=storage_dir,
env_manager=env_manager,
run_id=run_id,
parameters=parameters,
)
env_vars = get_run_env_vars(run_id, experiment_id)
env_vars.update(get_databricks_env_vars(mlflow.get_tracking_uri()))
mlflow_run_subprocess = _run_mlflow_run_cmd(mlflow_run_arr, env_vars)
return LocalSubmittedRun(run_id, mlflow_run_subprocess)
def _build_mlflow_run_cmd(
uri, entry_point, docker_args, storage_dir, env_manager, run_id, parameters
):
"""
Build and return an array containing an ``mlflow run`` command that can be invoked to locally
run the project at the specified URI.
"""
mlflow_run_arr = ["mlflow", "run", uri, "-e", entry_point, "--run-id", run_id]
if docker_args is not None:
for key, value in docker_args.items():
args = key if isinstance(value, bool) else f"{key}={value}"
mlflow_run_arr.extend(["--docker-args", args])
if storage_dir is not None:
mlflow_run_arr.extend(["--storage-dir", storage_dir])
mlflow_run_arr.extend(["--env-manager", env_manager])
for key, value in parameters.items():
mlflow_run_arr.extend(["-P", f"{key}={value}"])
return mlflow_run_arr
def _run_mlflow_run_cmd(mlflow_run_arr, env_map):
"""
Invoke ``mlflow run`` in a subprocess, which in turn runs the entry point in a child process.
Returns a handle to the subprocess. Popen launched to invoke ``mlflow run``.
"""
final_env = os.environ.copy()
final_env.update(env_map)
# Launch `mlflow run` command as the leader of its own process group so that we can do a
# best-effort cleanup of all its descendant processes if needed
if sys.platform == "win32":
return subprocess.Popen(
mlflow_run_arr,
env=final_env,
text=True,
creationflags=subprocess.CREATE_NEW_PROCESS_GROUP,
)
else:
return subprocess.Popen(mlflow_run_arr, env=final_env, text=True, preexec_fn=os.setsid)
def _run_entry_point(command, work_dir, experiment_id, run_id):
"""
Run an entry point command in a subprocess, returning a SubmittedRun that can be used to
query the run's status.
Args:
command: Entry point command to run
work_dir: Working directory in which to run the command
run_id: MLflow run ID associated with the entry point execution.
"""
env = os.environ.copy()
env.update(get_run_env_vars(run_id, experiment_id))
env.update(get_databricks_env_vars(tracking_uri=mlflow.get_tracking_uri()))
_logger.info("=== Running command '%s' in run with ID '%s' === ", command, run_id)
# in case os name is not 'nt', we are not running on windows. It introduces
# bash command otherwise.
if not is_windows():
process = subprocess.Popen(["bash", "-c", command], close_fds=True, cwd=work_dir, env=env)
else:
# process = subprocess.Popen(command, close_fds=True, cwd=work_dir, env=env)
process = subprocess.Popen(["cmd", "/c", command], close_fds=True, cwd=work_dir, env=env)
return LocalSubmittedRun(run_id, process)
def _get_docker_command(image, active_run, docker_args=None, volumes=None, user_env_vars=None):
from mlflow.projects.docker import get_docker_tracking_cmd_and_envs
docker_path = "docker"
cmd = [docker_path, "run", "--rm"]
if docker_args:
for name, value in docker_args.items():
# Passed just the name as boolean flag
if isinstance(value, bool) and value:
if len(name) == 1:
cmd += ["-" + name]
else:
cmd += ["--" + name]
else:
# Passed name=value
if len(name) == 1:
cmd += ["-" + name, value]
else:
cmd += ["--" + name, value]
env_vars = get_run_env_vars(
run_id=active_run.info.run_id, experiment_id=active_run.info.experiment_id
)
tracking_uri = tracking.get_tracking_uri()
tracking_cmds, tracking_envs = get_docker_tracking_cmd_and_envs(tracking_uri)
artifact_cmds, artifact_envs = _get_docker_artifact_storage_cmd_and_envs(
active_run.info.artifact_uri
)
cmd += tracking_cmds + artifact_cmds
env_vars.update(tracking_envs)
env_vars.update(artifact_envs)
if user_env_vars is not None:
for user_entry in user_env_vars:
if isinstance(user_entry, list):
# User has defined a new environment variable for the docker environment
env_vars[user_entry[0]] = user_entry[1]
else:
# User wants to copy an environment variable from system environment
system_var = os.environ.get(user_entry)
if system_var is None:
raise MlflowException(
"This project expects the {} environment variables to "
"be set on the machine running the project, but {} was "
"not set. Please ensure all expected environment variables "
"are set".format(", ".join(user_env_vars), user_entry)
)
env_vars[user_entry] = system_var
if volumes is not None:
for v in volumes:
cmd += ["-v", v]
for key, value in env_vars.items():
cmd += ["-e", f"{key}={value}"]
cmd += [image.tags[0]]
return cmd
def _get_local_artifact_cmd_and_envs(artifact_repo):
artifact_dir = artifact_repo.artifact_dir
container_path = artifact_dir
if not os.path.isabs(container_path):
container_path = os.path.join(MLFLOW_DOCKER_WORKDIR_PATH, container_path)
container_path = os.path.normpath(container_path)
abs_artifact_dir = os.path.abspath(artifact_dir)
return ["-v", f"{abs_artifact_dir}:{container_path}"], {}
def _get_s3_artifact_cmd_and_envs(artifact_repo):
if platform.system() == "Windows":
win_user_dir = os.environ["USERPROFILE"]
aws_path = os.path.join(win_user_dir, ".aws")
else:
aws_path = posixpath.expanduser("~/.aws")
volumes = []
if posixpath.exists(aws_path):
volumes = ["-v", "{}:{}".format(str(aws_path), "/.aws")]
envs = {
"AWS_SECRET_ACCESS_KEY": os.environ.get("AWS_SECRET_ACCESS_KEY"),
"AWS_ACCESS_KEY_ID": os.environ.get("AWS_ACCESS_KEY_ID"),
"MLFLOW_S3_ENDPOINT_URL": os.environ.get("MLFLOW_S3_ENDPOINT_URL"),
"MLFLOW_S3_IGNORE_TLS": os.environ.get("MLFLOW_S3_IGNORE_TLS"),
}
envs = {k: v for k, v in envs.items() if v is not None}
return volumes, envs
def _get_azure_blob_artifact_cmd_and_envs(artifact_repo):
envs = {
"AZURE_STORAGE_CONNECTION_STRING": os.environ.get("AZURE_STORAGE_CONNECTION_STRING"),
"AZURE_STORAGE_ACCESS_KEY": os.environ.get("AZURE_STORAGE_ACCESS_KEY"),
}
envs = {k: v for k, v in envs.items() if v is not None}
return [], envs
def _get_gcs_artifact_cmd_and_envs(artifact_repo):
cmds = []
envs = {}
if "GOOGLE_APPLICATION_CREDENTIALS" in os.environ:
credentials_path = os.environ["GOOGLE_APPLICATION_CREDENTIALS"]
cmds = ["-v", f"{credentials_path}:/.gcs"]
envs["GOOGLE_APPLICATION_CREDENTIALS"] = "/.gcs"
return cmds, envs
def _get_hdfs_artifact_cmd_and_envs(artifact_repo):
cmds = []
envs = {
"MLFLOW_KERBEROS_TICKET_CACHE": MLFLOW_KERBEROS_TICKET_CACHE.get(),
"MLFLOW_KERBEROS_USER": MLFLOW_KERBEROS_USER.get(),
"MLFLOW_PYARROW_EXTRA_CONF": MLFLOW_PYARROW_EXTRA_CONF.get(),
}
envs = {k: v for k, v in envs.items() if v is not None}
if "MLFLOW_KERBEROS_TICKET_CACHE" in envs:
ticket_cache = envs["MLFLOW_KERBEROS_TICKET_CACHE"]
cmds = ["-v", f"{ticket_cache}:{ticket_cache}"]
return cmds, envs
_artifact_storages = {
LocalArtifactRepository: _get_local_artifact_cmd_and_envs,
S3ArtifactRepository: _get_s3_artifact_cmd_and_envs,
AzureBlobArtifactRepository: _get_azure_blob_artifact_cmd_and_envs,
HdfsArtifactRepository: _get_hdfs_artifact_cmd_and_envs,
GCSArtifactRepository: _get_gcs_artifact_cmd_and_envs,
}
def _get_docker_artifact_storage_cmd_and_envs(artifact_uri):
artifact_repo = get_artifact_repository(artifact_uri)
_get_cmd_and_envs = _artifact_storages.get(type(artifact_repo))
if _get_cmd_and_envs is not None:
return _get_cmd_and_envs(artifact_repo)
else:
return [], {}