Repository: mlflow/mlflow

Issue:
[BUG] Spark model not recorded using `log_model` with distributed storage path
### Willingness to contribute
The MLflow Community encourages bug fix contributions. Would you or another member of your organization be willing to contribute a fix for this bug to the MLflow code base?

- [x] Yes. I can contribute a fix for this bug independently.
- [ ] Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.
- [ ] No. I cannot contribute a bug fix at this time.

### System information
- **Have I written custom code (as opposed to using a stock example script provided in MLflow)**: yes
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Linux PopOS 21.10
- **MLflow installed from (source or binary)**: binary
- **MLflow version (run ``mlflow --version``)**: 1.23.1, pyspark: 3.2.0
- **Python version**: 3.8
- **npm version, if running the dev UI**: N/A
- **Exact command to reproduce**: 

### Describe the problem
My experiment's `artifact_location` is a remote one (Google Storage). After training Spark model I'm using `mlflow.log_model` to register it - the model is saved into the correct location together with all the metadata, but it's not visible in the "runs list" view due to lack of tag `mlflow.log-model.history`. It's also treated as a normal (non-model) artifact on the run details view.

When using another experiment with the local artifacts path - everything works as expected.

### Code to reproduce issue

First, download GCS connector to use the direct way of logging the model:

```
wget https://repo1.maven.org/maven2/com/google/cloud/bigdataoss/gcs-connector/hadoop3-2.2.4/gcs-connector-hadoop3-2.2.4-shaded.jar
```

Then, run the following code:

```
import mlflow
from pyspark.ml import Pipeline
from pyspark.ml.classification import LogisticRegression
from pyspark.ml.feature import HashingTF, Tokenizer
from pyspark.sql import SparkSession

spark = (SparkSession.builder
    .config('spark.jars', 'gcs-connector-hadoop3-2.2.4-shaded.jar')
    .config('spark.hadoop.fs.gs.impl', 'com.google.cloud.hadoop.fs.gcs.GoogleHadoopFileSystem')
    .getOrCreate())
mlflow.create_experiment("experiment-with-non-local-artifact-location", "gs://bucket-name/path")
mlflow.set_experiment("experiment-with-non-local-artifact-location")

with mlflow.start_run():
    training = spark.createDataFrame([
        (0, "a b c d e spark", 1.0),
        (1, "b d", 0.0),
        (2, "spark f g h", 1.0),
        (3, "hadoop mapreduce", 0.0) ], ["id", "text", "label"])
    tokenizer = Tokenizer(inputCol="text", outputCol="words")
    hashingTF = HashingTF(inputCol=tokenizer.getOutputCol(), outputCol="features")
    lr = LogisticRegression(maxIter=10, regParam=0.001)
    pipeline = Pipeline(stages=[tokenizer, hashingTF, lr])
    model = pipeline.fit(training)
    mlflow.spark.log_model(model, "spark-model")
    
    current_tags = mlflow.get_run(mlflow.active_run().info.run_id).data.tags
    assert mlflow.utils.mlflow_tags.MLFLOW_LOGGED_MODELS in current_tags
```

Output:

```
(...)
Traceback (most recent call last):
  File "test.py", line 27, in <module>
    assert mlflow.utils.mlflow_tags.MLFLOW_LOGGED_MODELS in current_tags
AssertionError
```

And model is not properly registered in the UI:

![image](https://user-images.githubusercontent.com/4659069/152795602-5ba6c23f-7ffa-452b-a178-9c7bc0d56432.png)


### Other info / logs

I figured out the issue in the code already and I'm preparing a PR now.


### What component(s), interfaces, languages, and integrations does this bug affect?
Components 
- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [x] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [x] `area/tracking`: Tracking Service, tracking client APIs, autologging

Interface 
- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

Language 
- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

Integrations
- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations



Discussion:

Resolving patch:
diff --git a/mlflow/spark.py b/mlflow/spark.py
--- a/mlflow/spark.py
+++ b/mlflow/spark.py
@@ -255,6 +255,7 @@ def log_model(
             input_example=input_example,
         )
         mlflow.tracking.fluent.log_artifacts(tmp_model_metadata_dir, artifact_path)
+        mlflow.tracking.fluent._record_logged_model(mlflow_model)
         if registered_model_name is not None:
             mlflow.register_model(
                 "runs:/%s/%s" % (run_id, artifact_path),


Test patch:
diff --git a/tests/spark/test_spark_model_export.py b/tests/spark/test_spark_model_export.py
--- a/tests/spark/test_spark_model_export.py
+++ b/tests/spark/test_spark_model_export.py
@@ -652,6 +652,16 @@ def test_pyspark_version_is_logged_without_dev_suffix(spark_model_iris):
             assert any(x == f"pyspark=={unaffected_version}" for x in pip_deps)
 
 
+@pytest.mark.large
+def test_model_is_recorded_when_using_direct_save(spark_model_iris):
+    # Patch `is_local_uri` to enforce direct model serialization to DFS
+    with mock.patch("mlflow.spark.is_local_uri", return_value=False):
+        with mlflow.start_run():
+            sparkm.log_model(spark_model=spark_model_iris.model, artifact_path="model")
+            current_tags = mlflow.get_run(mlflow.active_run().info.run_id).data.tags
+            assert mlflow.utils.mlflow_tags.MLFLOW_LOGGED_MODELS in current_tags
+
+
 def test_shutil_copytree_without_file_permissions(tmpdir):
     src_dir = tmpdir.mkdir("src-dir")
     dst_dir = tmpdir.mkdir("dst-dir")