✅ Upgrade signature format for mlflow==2.10 compatibility

Galileo-Galilei · Jan 28, 2024 · e886799 · e886799
1 parent 4c07f1c
commit e886799
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/tests/framework/cli/test_cli_modelify.py b/tests/framework/cli/test_cli_modelify.py
@@ -393,8 +393,8 @@ def test_modelify_with_infer_signature_arg(
     loaded_model = mlflow.pyfunc.load_model(f"runs:/{new_run_id}/model")
 
     assert loaded_model.metadata.get_input_schema().to_dict() == [
-        {"name": "col_int", "type": "long"},
-        {"name": "col_str", "type": "string"},
+        {"name": "col_int", "type": "long", "required": True},
+        {"name": "col_str", "type": "string", "required": True},
     ]
 
 

diff --git a/tests/framework/hooks/test_hook_pipeline_ml.py b/tests/framework/hooks/test_hook_pipeline_ml.py
@@ -301,14 +301,14 @@ def test_mlflow_hook_save_pipeline_ml(
             assert (
                 (trained_model.metadata.signature.to_dict())
                 == {
-                    "inputs": '[{"name": "a", "type": "long"}]',
+                    "inputs": '[{"name": "a", "type": "long", "required": true}]',
                     "outputs": None,
                     "params": None,
                 }
             ) or (
                 (trained_model.metadata.signature.to_dict())
                 == {
-                    "inputs": '[{"type": "long", "name": "a"}]',
+                    "inputs": '[{"type": "long", "name": "a", "required": true}]',
                     "outputs": None,
                     "params": None,
                 }