mlflow#3975 Pass await_creation_for in run uri model registration (m…

…lflow#3976) Signed-off-by: Mark Greenwood <greenwma@googlemail.com> Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>
wamartin-aml · Jun 7, 2021 · 50f10ae · 50f10ae
1 parent 741e0da
commit 50f10ae
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git 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

diff --git 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,
         )