def test_mlflow_model_from_model_version(spark: SparkSession, mlflow_client):
# peg to a particular version of a model
spark.sql("CREATE MODEL resnet_m_fizz USING 'mlflow://rikai-test/1'")
check_ml_predict(spark, "resnet_m_fizz")
# use the latest version in a given stage (omitted means none)
spark.sql("CREATE MODEL resnet_m_buzz USING 'mlflow://rikai-test/'")
check_ml_predict(spark, "resnet_m_buzz")
def test_mlflow_model_from_runid(
spark: SparkSession, mlflow_client: MlflowClient
):
run_id = mlflow_client.search_model_versions("name='rikai-test'")[0].run_id
spark.sql(
"CREATE MODEL resnet_m_foo USING 'mlflow://{}/model'".format(run_id)
)
check_ml_predict(spark, "resnet_m_foo")
# if no path is given but only one artifact exists then use it by default
spark.sql("CREATE MODEL resnet_m_bar USING 'mlflow://{}'".format(run_id))
check_ml_predict(spark, "resnet_m_bar")
def test_mlflow_model_without_custom_logger(spark: SparkSession,
mlflow_client: MlflowClient):
spark.sql("CREATE MODEL vanilla_ice USING 'mlflow:/vanilla-mlflow/1'")
check_ml_predict(spark, "vanilla_ice")
schema = ("STRUCT<boxes:ARRAY<ARRAY<float>>,"
"scores:ARRAY<float>,labels:ARRAY<int>>")
pre_processing = ("rikai.contrib.torch.transforms."
"fasterrcnn_resnet50_fpn.pre_processing")
post_processing = ("rikai.contrib.torch.transforms."
"fasterrcnn_resnet50_fpn.post_processing")
spark.sql(("CREATE MODEL vanilla_fire "
"FLAVOR pytorch "
"PREPROCESSOR '{}' "
"POSTPROCESSOR '{}' "
"RETURNS {} "
"USING 'mlflow:/vanilla-mlflow-no-tags/1'").format(
pre_processing, post_processing, schema))
check_ml_predict(spark, "vanilla_fire")
def test_yaml_model(spark: SparkSession, resnet_spec: str):
spark.sql("CREATE MODEL resnet_m USING 'file://{}'".format(resnet_spec))
check_ml_predict(spark, "resnet_m")