Upgrade dependencies to handle MLFlow 2.x

README.md

@@ -69,20 +69,23 @@ deployment = client.create_deployment(
     config=dict(
         # Deployed model config
         machine_type="n1-standard-2",
-        min_replica_count=None,
-        max_replica_count=None,
+        min_replica_count=1,
+        max_replica_count=1,
         accelerator_type=None,
-        accelerator_count=None,
+        accelerator_count=0,
         service_account=None,
-        explanation_metadata=None, # JSON string
-        explanation_parameters=None, # JSON string
+        # explanation_metadata=None, # JSON string
+        # explanation_parameters=None, # JSON string
 
         # Model container image building config
         destination_image_uri=None,
 
+        # Model deployment config
+        sync="true",
+
         # Endpoint config
         endpoint_description=None,
-        endpoint_deploy_timeout=None,
+        # endpoint_deploy_timeout=1800,
 
         # Vertex AI config
         project=None,

google_cloud_mlflow/_mlflow_model_gcp_deployment_utils.py

@@ -61,20 +61,6 @@
 _logger = logging.getLogger(__name__)
 
 
-def get_fixed_mlflow_source_dir():
-    """Downloads the fixed MLflow source code."""
-    fixed_mlflow_archive_url = "https://github.com/Ark-kun/mlflow/archive/refs/heads/MLFlow-fixes.zip"
-    fixed_mlflow_archive_path, _ = urllib.request.urlretrieve(url=fixed_mlflow_archive_url)
-    fixed_mlflow_parent_dir = tempfile.mkdtemp(prefix="mlflow.fixed")
-    with zipfile.ZipFile(fixed_mlflow_archive_path, 'r') as zip_ref:
-        zip_ref.extractall(fixed_mlflow_parent_dir)
-    # The archive contains a subdirectory: "Ark-kun-mlflow-0ec4c64"
-    # So we need to go one level deeper
-    subdir = os.listdir(fixed_mlflow_parent_dir)[0]
-    fixed_mlflow_dir = os.path.join(fixed_mlflow_parent_dir, subdir)
-    return fixed_mlflow_dir
-
-
 def get_pickle_protocol(file_path: str) -> int:
     import pickletools
 
@@ -89,7 +75,7 @@ def get_pickle_protocol(file_path: str) -> int:
         except:
             pass
     return max_proto
-    
+
 
 def upload_mlflow_model_to_vertex_ai_models(
     model_uri: str,

google_cloud_mlflow/mlflow_model_deployment_plugin_for_google_cloud_vertex_ai.py

@@ -104,26 +104,29 @@ def create_deployment(
                 config=dict(
                     # Deployed model config
                     machine_type="n1-standard-2",
-                    min_replica_count=None,
-                    max_replica_count=None,
+                    min_replica_count=1,
+                    max_replica_count=1,
                     accelerator_type=None,
-                    accelerator_count=None,
-                    service_account=None,
-                    explanation_metadata=None, # JSON string
-                    explanation_parameters=None, # JSON string
+                    accelerator_count=0,
+                    # service_account=None,
+                    # explanation_metadata=None, # JSON string
+                    # explanation_parameters=None, # JSON string
 
                     # Model container image building config
                     destination_image_uri=None,
 
+                    # Model deployment config
+                    sync="true",
+
                     # Endpoint config
-                    endpoint_description=None,
-                    endpoint_deploy_timeout=None,
+                    endpoint_description="Model deployment",
+                    endpoint_deploy_timeout=1800,
 
                     # Vertex AI config
-                    project=None,
-                    location=None,
-                    encryption_spec_key_name=None,
-                    staging_bucket=None,
+                    # project=None,
+                    # location=None,
+                    # encryption_spec_key_name=None,
+                    # staging_bucket=None,
                 )
             )
 
@@ -217,7 +220,7 @@ def get_deployment(self, name: str) -> Dict[str, Any]:
         return _resource_to_mlflow_dict(
             self._get_deployment(deployment_name=name)
         )
-    
+
     def _get_deployment(self, deployment_name: str) -> aiplatform.Endpoint:
         endpoints = aiplatform.Endpoint.list(filter=f'display_name="{deployment_name}"')
         if len(endpoints) > 1:
@@ -422,13 +425,13 @@ def target_help():
                 config=dict(
                     # Deployed model config
                     machine_type="n1-standard-2",
-                    min_replica_count=None,
-                    max_replica_count=None,
+                    min_replica_count=1,
+                    max_replica_count=1,
                     accelerator_type=None,
-                    accelerator_count=None,
+                    accelerator_count=0,
                     service_account=None,
-                    explanation_metadata=None, # JSON string
-                    explanation_parameters=None, # JSON string
+                    # explanation_metadata=None, # JSON string
+                    # explanation_parameters=None, # JSON string
 
                     # Model container image building config
                     destination_image_uri=None,

setup.py

@@ -43,10 +43,10 @@
     keywords='mlflow, Google Cloud, Vertex AI',
     python_requires=">=3.6",
     install_requires=[
-        "google-cloud-aiplatform~=1.7",
-        "mlflow~=1.22",
-        "google-cloud-build==3.*",
-        "google-cloud-storage==1.*",
+        "google-cloud-aiplatform~=1.22",
+        "mlflow>=2",
+        "google-cloud-build~=3.13",
+        "google-cloud-storage~=2.7",
     ],
     entry_points={
         "mlflow.deployments": [