def test_change_conda_env_root_location(tmp_path, sk_model):
env_root1_path = tmp_path / "root1"
env_root1_path.mkdir()
env_root2_path = tmp_path / "root2"
env_root2_path.mkdir()
model1_path = tmp_path / "model1"
mlflow.sklearn.save_model(sk_model,
str(model1_path),
pip_requirements=["scikit-learn==1.0.1"])
model2_path = tmp_path / "model2"
mlflow.sklearn.save_model(sk_model,
str(model2_path),
pip_requirements=["scikit-learn==1.0.2"])
env_path_set = set()
for env_root_path, model_path, sklearn_ver in [
(env_root1_path, model1_path, "1.0.1"),
(
env_root2_path,
model1_path,
"1.0.1",
), # test the same env created in different env root path.
(
env_root1_path,
model2_path,
"1.0.2",
), # test different env created in the same env root path.
]:
_get_flavor_backend(
str(model_path),
env_manager=_EnvManager.CONDA,
install_mlflow=False,
env_root_dir=str(env_root_path),
).prepare_env(model_uri=str(model_path))
conda_env_name = _get_conda_env_name(str(model_path / "conda.yaml"),
env_root_dir=env_root_path)
env_path = env_root_path / "conda_envs" / conda_env_name
assert env_path.exists()
env_path_set.add(str(env_path))
python_exec_path = str(env_path / "bin" / "python")
# Test `_execute_in_conda_env` run command under the correct activated python env.
_execute_in_conda_env(
conda_env_name,
command=
f"python -c \"import sys; assert sys.executable == '{python_exec_path}'; "
f"import sklearn; assert sklearn.__version__ == '{sklearn_ver}'\"",
install_mlflow=False,
env_root_dir=str(env_root_path),
)
assert len(env_path_set) == 3
def test_spark_udf_embedded_model_server_killed_when_job_canceled(
spark, sklearn_model, model_path, env_manager):
from mlflow.pyfunc.scoring_server.client import ScoringServerClient
from mlflow.models.cli import _get_flavor_backend
mlflow.sklearn.save_model(sklearn_model.model, model_path)
server_port = 51234
timeout = 60
@pandas_udf("int")
def udf_with_model_server(it: Iterator[pd.Series]) -> Iterator[pd.Series]:
from mlflow.models.cli import _get_flavor_backend
_get_flavor_backend(model_path,
env_manager=env_manager,
workers=1,
install_mlflow=False).serve(
model_uri=model_path,
port=server_port,
host="127.0.0.1",
timeout=timeout,
enable_mlserver=False,
synchronous=False,
)
time.sleep(120)
for x in it:
yield x
def run_job():
# Start a spark job with only one UDF task,
# and the udf task starts a mlflow model server process.
spark.range(1).repartition(1).select(
udf_with_model_server("id")).collect()
_get_flavor_backend(model_path,
env_manager=env_manager,
install_mlflow=False).prepare_env(model_uri=model_path)
job_thread = threading.Thread(target=run_job)
job_thread.start()
client = ScoringServerClient("127.0.0.1", server_port)
client.wait_server_ready(timeout=20)
spark.sparkContext.cancelAllJobs()
job_thread.join()
time.sleep(10) # waiting server to exit and release the port.
# assert ping failed, i.e. the server process is killed successfully.
with pytest.raises(Exception): # pylint: disable=pytest-raises-without-match
client.ping()
def udf_with_model_server(it: Iterator[pd.Series]) -> Iterator[pd.Series]:
from mlflow.models.cli import _get_flavor_backend
_get_flavor_backend(
model_path, env_manager=_EnvManager.CONDA, workers=1, install_mlflow=False
).serve(
model_uri=model_path,
port=server_port,
host="127.0.0.1",
enable_mlserver=False,
synchronous=False,
)
time.sleep(120)
for x in it:
yield x
def test_scoring_server_client(sklearn_model, model_path):
from mlflow.pyfunc.scoring_server.client import ScoringServerClient
from mlflow.utils import find_free_port
from mlflow.models.cli import _get_flavor_backend
mlflow.sklearn.save_model(sk_model=sklearn_model.model, path=model_path)
expected_result = sklearn_model.model.predict(sklearn_model.inference_data)
port = find_free_port()
server_proc = None
try:
server_proc = _get_flavor_backend(model_path,
eng_manager=_EnvManager.CONDA,
workers=1,
install_mlflow=False).serve(
model_uri=model_path,
port=port,
host="127.0.0.1",
enable_mlserver=False,
synchronous=False,
)
client = ScoringServerClient(host="127.0.0.1", port=port)
client.wait_server_ready()
data = pd.DataFrame(sklearn_model.inference_data)
result = client.invoke(data).to_numpy()[:, 0]
np.testing.assert_allclose(result, expected_result, rtol=1e-5)
finally:
if server_proc is not None:
os.kill(server_proc.pid, signal.SIGTERM)
def _build_serving_image(model_uri: str,
destination_image_uri: str,
mlflow_source_dir: str = None):
_logger.info("Building image")
flavor_backend = _get_flavor_backend(model_uri)
flavor_backend.build_image(model_uri,
destination_image_uri,
install_mlflow=mlflow_source_dir is not None,
mlflow_home=mlflow_source_dir)
_logger.info("Uploading image to Google Container Registry")
client = docker.from_env()
result = client.images.push(destination_image_uri,
stream=True,
decode=True)
for line in result:
# Docker client doesn't catch auth errors, so we have to do it
# ourselves. See https://github.com/docker/docker-py/issues/1772
if 'errorDetail' in line:
raise docker.errors.APIError(line['errorDetail']['message'])
if 'status' in line:
print(line['status'])
