wehub-resource-sync
374 lines
15 KiB
Python
374 lines
15 KiB
Python
"""Internal utilities for parsing MLproject YAML files."""
|
|
|
|
import os
|
|
|
|
import yaml
|
|
|
|
from mlflow.exceptions import ExecutionException, MlflowException
|
|
from mlflow.projects import env_type
|
|
from mlflow.tracking import artifact_utils
|
|
from mlflow.utils import data_utils
|
|
from mlflow.utils.environment import _PYTHON_ENV_FILE_NAME
|
|
from mlflow.utils.file_utils import get_local_path_or_none
|
|
from mlflow.utils.string_utils import is_string_type, quote
|
|
|
|
MLPROJECT_FILE_NAME = "mlproject"
|
|
DEFAULT_CONDA_FILE_NAME = "conda.yaml"
|
|
|
|
|
|
def _find_mlproject(directory):
|
|
filenames = os.listdir(directory)
|
|
for filename in filenames:
|
|
if filename.lower() == MLPROJECT_FILE_NAME:
|
|
return os.path.join(directory, filename)
|
|
return None
|
|
|
|
|
|
def load_project(directory):
|
|
mlproject_path = _find_mlproject(directory)
|
|
|
|
# TODO: Validate structure of YAML loaded from the file
|
|
yaml_obj = {}
|
|
if mlproject_path is not None:
|
|
with open(mlproject_path) as mlproject_file:
|
|
yaml_obj = yaml.safe_load(mlproject_file)
|
|
|
|
# Validate the project config does't contain multiple environment fields
|
|
env_fields = set(yaml_obj.keys()).intersection(env_type.ALL)
|
|
if len(env_fields) > 1:
|
|
raise ExecutionException(
|
|
f"Project cannot contain multiple environment fields: {env_fields}"
|
|
)
|
|
|
|
project_name = yaml_obj.get("name")
|
|
|
|
# Parse entry points
|
|
entry_points = {}
|
|
for name, entry_point_yaml in yaml_obj.get("entry_points", {}).items():
|
|
parameters = entry_point_yaml.get("parameters", {})
|
|
command = entry_point_yaml.get("command")
|
|
entry_points[name] = EntryPoint(name, parameters, command)
|
|
|
|
databricks_spark_job_yaml = yaml_obj.get("databricks_spark_job")
|
|
if databricks_spark_job_yaml is not None:
|
|
python_file = databricks_spark_job_yaml.get("python_file")
|
|
|
|
if python_file is None and not entry_points:
|
|
raise MlflowException(
|
|
"Databricks Spark job requires either 'databricks_spark_job.python_file' "
|
|
"setting or 'entry_points' setting."
|
|
)
|
|
if python_file is not None and entry_points:
|
|
raise MlflowException(
|
|
"Databricks Spark job does not allow setting both "
|
|
"'databricks_spark_job.python_file' and 'entry_points'."
|
|
)
|
|
|
|
for entry_point in entry_points.values():
|
|
for param in entry_point.parameters.values():
|
|
if param.type == "path":
|
|
raise MlflowException(
|
|
"Databricks Spark job does not support entry point parameter of 'path' "
|
|
f"type. '{param.name}' value type is invalid."
|
|
)
|
|
|
|
if env_type.DOCKER in yaml_obj:
|
|
raise MlflowException(
|
|
"Databricks Spark job does not support setting docker environment."
|
|
)
|
|
|
|
if env_type.PYTHON in yaml_obj:
|
|
raise MlflowException(
|
|
"Databricks Spark job does not support setting python environment."
|
|
)
|
|
|
|
if env_type.CONDA in yaml_obj:
|
|
raise MlflowException(
|
|
"Databricks Spark job does not support setting conda environment."
|
|
)
|
|
|
|
databricks_spark_job_spec = DatabricksSparkJobSpec(
|
|
python_file=databricks_spark_job_yaml.get("python_file"),
|
|
parameters=databricks_spark_job_yaml.get("parameters", []),
|
|
python_libraries=databricks_spark_job_yaml.get("python_libraries", []),
|
|
)
|
|
return Project(
|
|
databricks_spark_job_spec=databricks_spark_job_spec,
|
|
name=project_name,
|
|
entry_points=entry_points,
|
|
)
|
|
|
|
# Validate config if docker_env parameter is present
|
|
if docker_env := yaml_obj.get(env_type.DOCKER):
|
|
if not docker_env.get("image"):
|
|
raise ExecutionException(
|
|
"Project configuration (MLproject file) was invalid: Docker "
|
|
"environment specified but no image attribute found."
|
|
)
|
|
if docker_env.get("volumes"):
|
|
if not (
|
|
isinstance(docker_env["volumes"], list)
|
|
and all(isinstance(i, str) for i in docker_env["volumes"])
|
|
):
|
|
raise ExecutionException(
|
|
"Project configuration (MLproject file) was invalid: "
|
|
"Docker volumes must be a list of strings, "
|
|
"""e.g.: '["/path1/:/path1", "/path2/:/path2"])"""
|
|
)
|
|
if docker_env.get("environment"):
|
|
if not (
|
|
isinstance(docker_env["environment"], list)
|
|
and all(isinstance(i, (list, str)) for i in docker_env["environment"])
|
|
):
|
|
raise ExecutionException(
|
|
"Project configuration (MLproject file) was invalid: "
|
|
"environment must be a list containing either strings (to copy environment "
|
|
"variables from host system) or lists of string pairs (to define new "
|
|
"environment variables)."
|
|
"""E.g.: '[["NEW_VAR", "new_value"], "VAR_TO_COPY_FROM_HOST"])"""
|
|
)
|
|
return Project(
|
|
env_type=env_type.DOCKER,
|
|
env_config_path=None,
|
|
entry_points=entry_points,
|
|
docker_env=docker_env,
|
|
name=project_name,
|
|
)
|
|
|
|
if python_env := yaml_obj.get(env_type.PYTHON):
|
|
python_env_path = os.path.join(directory, python_env)
|
|
if not os.path.exists(python_env_path):
|
|
raise ExecutionException(
|
|
f"Project specified python_env file {python_env_path}, but no such file was found."
|
|
)
|
|
return Project(
|
|
env_type=env_type.PYTHON,
|
|
env_config_path=python_env_path,
|
|
entry_points=entry_points,
|
|
docker_env=None,
|
|
name=project_name,
|
|
)
|
|
|
|
if conda_path := yaml_obj.get(env_type.CONDA):
|
|
conda_env_path = os.path.join(directory, conda_path)
|
|
if not os.path.exists(conda_env_path):
|
|
raise ExecutionException(
|
|
f"Project specified conda environment file {conda_env_path}, but no such "
|
|
"file was found."
|
|
)
|
|
return Project(
|
|
env_type=env_type.CONDA,
|
|
env_config_path=conda_env_path,
|
|
entry_points=entry_points,
|
|
docker_env=None,
|
|
name=project_name,
|
|
)
|
|
|
|
default_python_env_path = os.path.join(directory, _PYTHON_ENV_FILE_NAME)
|
|
if os.path.exists(default_python_env_path):
|
|
return Project(
|
|
env_type=env_type.PYTHON,
|
|
env_config_path=default_python_env_path,
|
|
entry_points=entry_points,
|
|
docker_env=None,
|
|
name=project_name,
|
|
)
|
|
|
|
default_conda_path = os.path.join(directory, DEFAULT_CONDA_FILE_NAME)
|
|
if os.path.exists(default_conda_path):
|
|
return Project(
|
|
env_type=env_type.CONDA,
|
|
env_config_path=default_conda_path,
|
|
entry_points=entry_points,
|
|
docker_env=None,
|
|
name=project_name,
|
|
)
|
|
|
|
return Project(
|
|
env_type=env_type.PYTHON,
|
|
env_config_path=None,
|
|
entry_points=entry_points,
|
|
docker_env=None,
|
|
name=project_name,
|
|
)
|
|
|
|
|
|
class Project:
|
|
"""A project specification loaded from an MLproject file in the passed-in directory."""
|
|
|
|
def __init__(
|
|
self,
|
|
name,
|
|
env_type=None,
|
|
env_config_path=None,
|
|
entry_points=None,
|
|
docker_env=None,
|
|
databricks_spark_job_spec=None,
|
|
):
|
|
self.env_type = env_type
|
|
self.env_config_path = env_config_path
|
|
self._entry_points = entry_points
|
|
self.docker_env = docker_env
|
|
self.name = name
|
|
self.databricks_spark_job_spec = databricks_spark_job_spec
|
|
|
|
def get_entry_point(self, entry_point):
|
|
if self.databricks_spark_job_spec:
|
|
if self.databricks_spark_job_spec.python_file is not None:
|
|
# If Databricks Spark job is configured with python_file field,
|
|
# it does not need to configure entry_point section
|
|
# and the 'entry_point' param in 'mlflow run' command is ignored
|
|
return None
|
|
|
|
if self._entry_points is None or entry_point not in self._entry_points:
|
|
raise MlflowException(
|
|
f"The entry point '{entry_point}' is not defined in the Databricks spark job "
|
|
f"MLproject file."
|
|
)
|
|
|
|
if entry_point in self._entry_points:
|
|
return self._entry_points[entry_point]
|
|
_, file_extension = os.path.splitext(entry_point)
|
|
ext_to_cmd = {".py": "python", ".sh": os.environ.get("SHELL", "bash")}
|
|
if file_extension in ext_to_cmd:
|
|
command = f"{ext_to_cmd[file_extension]} {quote(entry_point)}"
|
|
if not is_string_type(command):
|
|
command = command.encode("utf-8")
|
|
return EntryPoint(name=entry_point, parameters={}, command=command)
|
|
elif file_extension == ".R":
|
|
command = f"Rscript -e \"mlflow::mlflow_source('{quote(entry_point)}')\" --args"
|
|
return EntryPoint(name=entry_point, parameters={}, command=command)
|
|
raise ExecutionException(
|
|
"Could not find {0} among entry points {1} or interpret {0} as a "
|
|
"runnable script. Supported script file extensions: "
|
|
"{2}".format(entry_point, list(self._entry_points.keys()), list(ext_to_cmd.keys()))
|
|
)
|
|
|
|
|
|
class EntryPoint:
|
|
"""An entry point in an MLproject specification."""
|
|
|
|
def __init__(self, name, parameters, command):
|
|
self.name = name
|
|
self.parameters = {k: Parameter(k, v) for (k, v) in parameters.items()}
|
|
self.command = command
|
|
|
|
def _validate_parameters(self, user_parameters):
|
|
missing_params = [
|
|
name
|
|
for name in self.parameters
|
|
if name not in user_parameters and self.parameters[name].default is None
|
|
]
|
|
if missing_params:
|
|
raise ExecutionException(
|
|
"No value given for missing parameters: {}".format(
|
|
", ".join([f"'{name}'" for name in missing_params])
|
|
)
|
|
)
|
|
|
|
def compute_parameters(self, user_parameters, storage_dir):
|
|
"""
|
|
Given a dict mapping user-specified param names to values, computes parameters to
|
|
substitute into the command for this entry point. Returns a tuple (params, extra_params)
|
|
where `params` contains key-value pairs for parameters specified in the entry point
|
|
definition, and `extra_params` contains key-value pairs for additional parameters passed
|
|
by the user.
|
|
|
|
Note that resolving parameter values can be a heavy operation, e.g. if a remote URI is
|
|
passed for a parameter of type `path`, we download the URI to a local path within
|
|
`storage_dir` and substitute in the local path as the parameter value.
|
|
|
|
If `storage_dir` is `None`, report path will be return as parameter.
|
|
"""
|
|
if user_parameters is None:
|
|
user_parameters = {}
|
|
# Validate params before attempting to resolve parameter values
|
|
self._validate_parameters(user_parameters)
|
|
final_params = {}
|
|
extra_params = {}
|
|
|
|
parameter_keys = list(self.parameters.keys())
|
|
for key in parameter_keys:
|
|
param_obj = self.parameters[key]
|
|
key_position = parameter_keys.index(key)
|
|
value = user_parameters[key] if key in user_parameters else self.parameters[key].default
|
|
final_params[key] = param_obj.compute_value(value, storage_dir, key_position)
|
|
for key in user_parameters:
|
|
if key not in final_params:
|
|
extra_params[key] = user_parameters[key]
|
|
return (
|
|
self._sanitize_value_dict(final_params),
|
|
self._sanitize_extra_param_dict(extra_params),
|
|
)
|
|
|
|
def compute_command(self, user_parameters, storage_dir):
|
|
params, extra_params = self.compute_parameters(user_parameters, storage_dir)
|
|
command_with_params = self.command.format(**params)
|
|
command_arr = [command_with_params]
|
|
command_arr.extend([f"--{key} {value}" for key, value in extra_params.items()])
|
|
return " ".join(command_arr)
|
|
|
|
@staticmethod
|
|
def _sanitize_value_dict(param_dict):
|
|
# Keys here are used as str.format placeholders against self.command,
|
|
# not shell tokens, so quoting them would break {placeholder} resolution.
|
|
# Only values flow into the shell command unquoted, so only values need
|
|
# shell-quoting.
|
|
return {str(key): quote(str(value)) for key, value in param_dict.items()}
|
|
|
|
@staticmethod
|
|
def _sanitize_extra_param_dict(param_dict):
|
|
# Both keys and values get joined into the shell command as
|
|
# `--{key} {value}` tokens, so both need shell-quoting.
|
|
return {quote(str(key)): quote(str(value)) for key, value in param_dict.items()}
|
|
|
|
|
|
class Parameter:
|
|
"""A parameter in an MLproject entry point."""
|
|
|
|
def __init__(self, name, yaml_obj):
|
|
self.name = name
|
|
if is_string_type(yaml_obj):
|
|
self.type = yaml_obj
|
|
self.default = None
|
|
else:
|
|
self.type = yaml_obj.get("type", "string")
|
|
self.default = yaml_obj.get("default")
|
|
|
|
def _compute_uri_value(self, user_param_value):
|
|
if not data_utils.is_uri(user_param_value):
|
|
raise ExecutionException(
|
|
f"Expected URI for parameter {self.name} but got {user_param_value}"
|
|
)
|
|
return user_param_value
|
|
|
|
def _compute_path_value(self, user_param_value, storage_dir, key_position):
|
|
if local_path := get_local_path_or_none(user_param_value):
|
|
if not os.path.exists(local_path):
|
|
raise ExecutionException(
|
|
f"Got value {user_param_value} for parameter {self.name}, but no such file or "
|
|
"directory was found."
|
|
)
|
|
return os.path.abspath(local_path)
|
|
target_sub_dir = f"param_{key_position}"
|
|
download_dir = os.path.join(storage_dir, target_sub_dir)
|
|
os.mkdir(download_dir)
|
|
return artifact_utils._download_artifact_from_uri(
|
|
artifact_uri=user_param_value, output_path=download_dir
|
|
)
|
|
|
|
def compute_value(self, param_value, storage_dir, key_position):
|
|
if storage_dir and self.type == "path":
|
|
return self._compute_path_value(param_value, storage_dir, key_position)
|
|
elif self.type == "uri":
|
|
return self._compute_uri_value(param_value)
|
|
else:
|
|
return param_value
|
|
|
|
|
|
class DatabricksSparkJobSpec:
|
|
def __init__(self, python_file, parameters, python_libraries):
|
|
self.python_file = python_file
|
|
self.parameters = parameters
|
|
self.python_libraries = python_libraries
|