def test_enabling_autologging_before_spark_session_works(disable):
mlflow.spark.autolog(disable=disable)
# creating spark session AFTER autolog was enabled
spark_session = _get_or_create_spark_session()
rows = [Row(100)]
schema = StructType([StructField("number2", IntegerType())])
rdd = spark_session.sparkContext.parallelize(rows)
df = spark_session.createDataFrame(rdd, schema)
tempdir = tempfile.mkdtemp()
filepath = os.path.join(tempdir, "test-data")
df.write.option("header", "true").format("csv").save(filepath)
read_df = (spark_session.read.format("csv").option(
"header", "true").option("inferSchema", "true").load(filepath))
with mlflow.start_run():
run_id = mlflow.active_run().info.run_id
read_df.collect()
time.sleep(1)
run = mlflow.get_run(run_id)
if disable:
_assert_spark_data_not_logged(run=run)
else:
_assert_spark_data_logged(run=run, path=filepath, data_format="csv")
shutil.rmtree(tempdir)
spark_session.stop()
def test_enabling_autologging_throws_for_missing_jar():
# pylint: disable=unused-argument
spark_session = _get_or_create_spark_session(jars="")
try:
with pytest.raises(MlflowException) as exc:
kiwi.spark.autolog()
assert "Please ensure you have the mlflow-spark JAR attached" in exc.value.message
finally:
spark_session.stop()
def spark_session():
session = _get_or_create_spark_session()
yield session
session.stop()