Repository: mlflow/mlflow

Issue:
[BUG] await_registration_for isn't always overridable when registering model version
### 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?

- [ ] Yes. I can contribute a fix for this bug independently.
- [x] 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)**: No
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Databricks
- **MLflow installed from (source or binary)**: NA
- **MLflow version (run ``mlflow --version``)**: 1.13.1
- **Python version**: 3.6
- **npm version, if running the dev UI**: NA
- **Exact command to reproduce**: mlflow.spark.log_model(spark_model=model, registered_model_name = "test-model", artifact_path="test-model", await_registration_for=600)

### Describe the problem
Apologies if this is desired behaviour, but passing await_registration_for didn't seem to have an effect in the case above for us - the default wait of 300 seconds was always used.  Looking in to the mlflow/tracking/_model_registry/fluent.py it seems like [L69 doesn't pass this forward to client.create_model down a certain route](https://github.com/mlflow/mlflow/blob/75c4bfa1c0dda546c491adfed670f8a80b9f3ccf/mlflow/tracking/_model_registry/fluent.py#L69).

If there is a case where this shouldn't be allowed it isn't called out in the documentation.

### Code to reproduce issue
mlflow.start_run()
mlflow.spark.log_model(spark_model=model, registered_model_name = "test-model", artifact_path="test-model", await_registration_for=600)

### Other info / logs
2021/01/12 16:43:18 INFO mlflow.tracking._model_registry.client: Waiting up to 300 seconds for model version to finish creation. Model name: test-model, version 1

### 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
- [x] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: Local serving, model deployment tools, spark UDFs
- [ ] `area/server-infra`: MLflow server, JavaScript dev server
- [x] `area/tracking`: Tracking Service, tracking client APIs, autologging

Interface 
- [ ] `area/uiux`: Front-end, user experience, JavaScript, plotting
- [ ] `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:
@whiteh Thanks for reporthing this. This looks like a bug to me.

```python
# test.py

from unittest import mock

import mlflow


def test_create_model_version_called_with_await_creation_for():
    mlflow.set_tracking_uri("sqlite:///mlflow.db")
    mock_target = "mlflow.tracking.client.MlflowClient.create_model_version"

    with mock.patch(mock_target) as mock_create_model_version, mlflow.start_run() as run:
        model = "mock_model"  # HACK: `mlflow.sklearn.log_model` accepts anything picklable
        artifact_path = "model"
        mlflow.sklearn.log_model(model, artifact_path)
        model_uri = mlflow.get_artifact_uri(artifact_path)

        registered_model_name = "model"
        model_uri_in_runs_format = f"runs:/{run.info.run_id}/{artifact_path}"
        model_uri = mlflow.get_artifact_uri(artifact_path)

        mlflow.register_model(
            model_uri_in_runs_format, registered_model_name, await_registration_for=100
        )
        # `create_model_version` should have been called with `await_creation_for=100`
        mock_create_model_version.assert_called_once_with(
            registered_model_name, model_uri, run.info.run_id, await_creation_for=100
        )
```

This gives:

```
>>> pytest test.py

...
E           AssertionError:
E           Expected call: create_model_version('model', './mlruns/0/8d4c08a11ee04514a328b093ee98b816/artifacts/model', '8d4c08a11ee04514a328b093ee98b816', await_creation_for=100)
E           Actual call:   create_model_version('model', './mlruns/0/8d4c08a11ee04514a328b093ee98b816/artifacts/model', '8d4c08a11ee04514a328b093ee98b816')
```


cc @andychow-db @sueann 

Resolving patch:
diff --git a/mlflow/tracking/_model_registry/fluent.py b/mlflow/tracking/_model_registry/fluent.py
--- a/mlflow/tracking/_model_registry/fluent.py
+++ b/mlflow/tracking/_model_registry/fluent.py
@@ -66,7 +66,9 @@ def register_model(model_uri, name, await_registration_for=DEFAULT_AWAIT_MAX_SLE
     if RunsArtifactRepository.is_runs_uri(model_uri):
         source = RunsArtifactRepository.get_underlying_uri(model_uri)
         (run_id, _) = RunsArtifactRepository.parse_runs_uri(model_uri)
-        create_version_response = client.create_model_version(name, source, run_id)
+        create_version_response = client.create_model_version(
+            name, source, run_id, await_creation_for=await_registration_for
+        )
     else:
         create_version_response = client.create_model_version(
             name, source=model_uri, run_id=None, await_creation_for=await_registration_for


Test patch:
diff --git a/tests/tracking/_model_registry/test_model_registry_fluent.py b/tests/tracking/_model_registry/test_model_registry_fluent.py
--- a/tests/tracking/_model_registry/test_model_registry_fluent.py
+++ b/tests/tracking/_model_registry/test_model_registry_fluent.py
@@ -49,7 +49,10 @@ def test_register_model_with_runs_uri():
         register_model("runs:/run12345/path/to/model", "Model 1")
         MlflowClient.create_registered_model.assert_called_once_with("Model 1")
         MlflowClient.create_model_version.assert_called_once_with(
-            "Model 1", "s3:/path/to/source", "run12345"
+            "Model 1",
+            "s3:/path/to/source",
+            "run12345",
+            await_creation_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
         )