mlflow-tclake-plugin 2.0.1__py3-none-any.whl → 3.0.1__py3-none-any.whl

mlflow_tclake_plugin/tclake_store.py

@@ -3,6 +3,7 @@ import json
 import os
 import time
 import uuid
+import urllib.parse
 from datetime import datetime
 
 from cachetools import TTLCache
@@ -10,16 +11,23 @@ from tencentcloud.common import credential
 from tencentcloud.common.common_client import CommonClient
 from tencentcloud.common.profile.client_profile import ClientProfile
 
-from mlflow.entities.model_registry import RegisteredModel, ModelVersion, ModelVersionTag
+from mlflow.entities.model_registry import (
+    RegisteredModel,
+    ModelVersion,
+    ModelVersionTag,
+)
 from mlflow.exceptions import MlflowException
-from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
-from mlflow.store.model_registry import SEARCH_REGISTERED_MODEL_MAX_RESULTS_THRESHOLD, \
-    SEARCH_MODEL_VERSION_MAX_RESULTS_THRESHOLD
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE, RESOURCE_ALREADY_EXISTS
+from mlflow.store.model_registry import (
+    SEARCH_REGISTERED_MODEL_MAX_RESULTS_THRESHOLD,
+    SEARCH_MODEL_VERSION_MAX_RESULTS_THRESHOLD,
+)
 from mlflow.store.model_registry.abstract_store import AbstractStore
 from mlflow.utils.annotations import experimental
 from mlflow.utils.search_utils import SearchModelUtils, SearchModelVersionUtils
 from mlflow.store.entities.paged_list import PagedList
-
+from mlflow.models.model import get_model_info
+from tencentcloud.common.exception.tencent_cloud_sdk_exception import TencentCloudSDKException
 
 def to_string(obj):
     if obj is None:
@@ -65,9 +73,13 @@ def _get_description(resp):
 TCLAKE_MLFLOW_TAG_PREFIX = "tclake.tag."
 TCLAKE_MLFLOW_RUN_ID_KEY = "tclake.mlflow.run_id"
 TCLAKE_MLFLOW_RUN_LINK_KEY = "tclake.mlflow.run_link"
-TCLAKE_UUID_KEY = "tccatalog.identifier"
+TCLAKE_MLFLOW_DEPLOYMENT_JOB_ID_KEY = "tclake.mlflow.deployment_job_id"
+TCLAKE_WEDATA_WORKSPACE_ID_KEY = "tclake.wedata.workspace_id"
+TCLAKE_MLFLOW_DEPLOYMENT_STATUS_KEY = "tclake.mlflow.deployment_status"
+TCLAKE_MLFLOW_MODEL_ID_KEY = "tclake.mlflow.model_id"
 TCLAKE_MLFLOW_MODEL_SIGNATURE_KEY = "tclake.mlflow.model_signature"
 
+UN_DEPLOYMENT = "UnDeployment"
 
 def _set_kv_to_properties(key, value, properties=None):
     if properties is None:
@@ -86,10 +98,6 @@ def _get_kv_from_properties(properties, key):
     return None
 
 
-def _get_uuid_from_properties(properties):
-    return _get_kv_from_properties(properties, TCLAKE_UUID_KEY)
-
-
 def _set_run_id_to_properties(run_id, properties=None):
     return _set_kv_to_properties(TCLAKE_MLFLOW_RUN_ID_KEY, run_id, properties)
 
@@ -105,10 +113,15 @@ def _set_run_link_to_properties(run_link, properties):
 def _get_run_link_from_properties(properties):
     return _get_kv_from_properties(properties, TCLAKE_MLFLOW_RUN_LINK_KEY)
 
+def _get_model_id_from_properties(properties):
+    return _get_kv_from_properties(properties, TCLAKE_MLFLOW_MODEL_ID_KEY)
+
 
 def _add_tag_to_properties(tag, properties=None):
     if tag and tag.value is not None:
-        properties = _set_kv_to_properties(TCLAKE_MLFLOW_TAG_PREFIX + tag.key, tag.value, properties)
+        properties = _set_kv_to_properties(
+            TCLAKE_MLFLOW_TAG_PREFIX + tag.key, tag.value, properties
+        )
     return properties
 
 
@@ -118,84 +131,47 @@ def _add_tags_to_properties(tags, properties=None):
             properties = _add_tag_to_properties(tag, properties)
     return properties
 
+def _add_deployment_job_id_to_properties(deployment_job_id, properties=None):
+    properties = _set_kv_to_properties(TCLAKE_MLFLOW_DEPLOYMENT_JOB_ID_KEY, deployment_job_id, properties)
+    return properties
+
+def _add_workspace_id_to_properties(workspace_id, properties=None):
+    if workspace_id is not None:
+        properties = _set_kv_to_properties(TCLAKE_WEDATA_WORKSPACE_ID_KEY, workspace_id, properties)
+    return properties
 
+def _add_deployment_status_to_properties(status, properties=None):
+    if status is not None:
+        properties = _set_kv_to_properties(TCLAKE_MLFLOW_DEPLOYMENT_STATUS_KEY, status, properties)
+    return properties
 
-def parse_model_signatures_from_dict(signature_dict):
-    signatures = []
-    inputs_str = signature_dict.get("inputs")
-    if inputs_str:
-        inputs = json.loads(inputs_str)
-        for input_item in inputs:
-            model_signature = {
-                "name": input_item.get("name", ""), 
-                "type": "INPUT",
-                "inputFlag": "true"
-            }
-            type_val = input_item.get("type")
-            if type_val is not None:
-                if type_val == "tensor":
-                    tensor_spec = input_item.get("tensor-spec", {}).get("dtype", "")
-                    if tensor_spec:
-                        model_signature["type"] = str(tensor_spec)
-                else:
-                    model_signature["type"] = str(type_val)
-            
-            signatures.append(model_signature)
-
-    outputs_str = signature_dict.get("outputs")
-    if outputs_str:
-        outputs = json.loads(outputs_str)
-        for output_item in outputs:    
-            name_val = output_item.get("name", "")
-            if not name_val:
-                name_val = output_item.get("prediction_column_name", "")
-            model_signature = {
-                "name": name_val, 
-                "type": "OUTPUT", 
-                "inputFlag": "false"
-            }
-            type_val = output_item.get("type")
-            if type_val is not None:
-                if type_val == "tensor":
-                    tensor_spec = output_item.get("tensor-spec", {}).get("dtype", "")
-                    if tensor_spec:
-                        model_signature["type"] = str(tensor_spec)
-                else:
-                    model_signature["type"] = str(type_val)
-            
-            signatures.append(model_signature)
-    
-    log_msg(f"parse model version signatures: {signatures}")
-    return signatures      
+def _add_model_id_to_properties(model_id, properties=None):
+    if model_id is not None:
+        properties = _set_kv_to_properties(TCLAKE_MLFLOW_MODEL_ID_KEY, model_id, properties)
+    return properties
 
 
 def _add_model_signature_to_properties(source, properties):
     log_msg("model source is {}".format(source))
-    try:
-        from mlflow.models.model import get_model_info
-        model_info = get_model_info(source)
-        signature = model_info.signature
-        log_msg("model {} signature is {}".format(source, signature))
-        if signature:
-            sig_json = json.dumps(parse_model_signatures_from_dict(signature.to_dict()))
-            log_msg("model {} signature json is {}".format(source, sig_json))
-        else:
-            log_msg(f"Registered model signature is not found in source artifact location '{source}'")
-            sig_json = json.dumps([])
-    except Exception as e:
-        log_msg(f"Failed to get model signature from {source}: {e}")
-        sig_json = json.dumps([])
+    model_info = get_model_info(source)  
+    signature = model_info.signature
+    log_msg("model {} signature is {}".format(source, signature))
+    if signature is None:
+        raise MlflowException(f"Registered model signature is not found in source artifact location '{source}'") 
+    sig_json = json.dumps(signature.to_dict()) 
+    log_msg("model {} signature json is {}".format(source, sig_json))
     properties = _set_kv_to_properties(TCLAKE_MLFLOW_MODEL_SIGNATURE_KEY, sig_json, properties)
     return properties
 
-
 def _get_tags_from_properties(properties):
     if properties is None:
         return None
     tags = []
     for p in properties:
         if p["Key"].startswith(TCLAKE_MLFLOW_TAG_PREFIX):
-            tags.append(ModelVersionTag(p["Key"][len(TCLAKE_MLFLOW_TAG_PREFIX):], p["Value"]))
+            tags.append(
+                ModelVersionTag(p["Key"][len(TCLAKE_MLFLOW_TAG_PREFIX) :], p["Value"])
+            )
     return tags
 
 
@@ -219,17 +195,20 @@ def _make_model(resp):
         creation_timestamp=_get_create_time_from_audit(audit),
         last_updated_timestamp=_get_last_updated_time_from_audit(audit),
         description=_get_description(resp),
-        tags=_get_tags_from_properties(properties),
+        tags=_get_tags_from_properties(properties)
     )
 
 
 def _get_model_version_name(entity):
-    return "{}.{}.{}".format(entity["CatalogName"], entity["SchemaName"], entity["ModelName"])
+    return "{}.{}.{}".format(
+        entity["CatalogName"], entity["SchemaName"], entity["ModelName"]
+    )
 
 
 def _make_model_version(entity, name):
     properties = entity["Properties"]
     audit = entity["Audit"]
+    aliases = entity["Aliases"]
     return ModelVersion(
         name=name,
         version=_get_model_version(entity["Version"]),
@@ -240,6 +219,8 @@ def _make_model_version(entity, name):
         run_id=_get_run_id_from_properties(properties),
         tags=_get_tags_from_properties(properties),
         run_link=_get_run_link_from_properties(properties),
+        aliases=aliases,
+        model_id=_get_model_id_from_properties(properties),
         status="READY",
     )
 
@@ -267,7 +248,9 @@ def _parse_page_token(page_token):
 
 
 def _create_page_token(offset, search_id):
-    return base64.b64encode(json.dumps({"offset": offset, "search_id": search_id}).encode("utf-8"))
+    return base64.b64encode(
+        json.dumps({"offset": offset, "search_id": search_id}).encode("utf-8")
+    )
 
 
 @experimental
@@ -278,29 +261,50 @@ class TCLakeStore(AbstractStore):
 
     def __init__(self, store_uri=None, tracking_uri=None):
         super().__init__(store_uri, tracking_uri)
-        log_msg("initializing tencent tclake client {} {}".format(store_uri, tracking_uri))
+        log_msg(
+            "initializing tencent tclake client {} {}".format(store_uri, tracking_uri)
+        )
         sid = os.getenv("TENCENTCLOUD_SECRET_ID", "")
         if len(sid) == 0:
             raise MlflowException("TENCENTCLOUD_SECRET_ID is not set")
         sk = os.getenv("TENCENTCLOUD_SECRET_KEY", "")
         if len(sk) == 0:
             raise MlflowException("TENCENTCLOUD_SECRET_KEY is not set")
+        sk = os.getenv("TENCENTCLOUD_SECRET_KEY", "")
+        if len(sk) == 0:
+            raise MlflowException("TENCENTCLOUD_SECRET_KEY is not set")
+
         token = os.getenv("TENCENTCLOUD_TOKEN", None)
+
+        self.workspace_id = os.getenv("WEDATA_WORKSPACE_ID", "")
+        if len(self.workspace_id) == 0:
+            raise MlflowException("WEDATA_WORKSPACE_ID is not set")
+        log_msg(str.format("wedata workspace id: {}", self.workspace_id))
+        
         client_profile = get_tencent_cloud_client_profile()
         cred = credential.Credential(sid, sk, token)
         parts = store_uri.split(":")
         if len(parts) < 2:
             raise MlflowException("set store_uri tclake:{region}")
         region = parts[1]
-        self.client = CommonClient("tccatalog", "2024-10-24", cred, region, client_profile)
+        self.client = CommonClient(
+            "tccatalog", "2024-10-24", cred, region, client_profile
+        )
         self.headers = get_tencent_cloud_headers()
-        self.default_catalog_name = os.getenv("TENCENTCLOUD_DEFAULT_CATALOG_NAME", "default")
-        self.default_schema_name = os.getenv("TENCENTCLOUD_DEFAULT_SCHEMA_NAME", "default")
+        self.default_catalog_name = os.getenv(
+            "TENCENTCLOUD_DEFAULT_CATALOG_NAME", "default"
+        )
+        self.default_schema_name = os.getenv(
+            "TENCENTCLOUD_DEFAULT_SCHEMA_NAME", "default"
+        )
         cache_size = int(os.getenv("TCLAKE_CACHE_SIZE", "100"))
         cache_ttl = int(os.getenv("TCLAKE_CACHE_TTL_SECS", "300"))
         self.cache = TTLCache(maxsize=cache_size, ttl=cache_ttl)
-        log_msg("initialized tencent tclake client successfully {} {} {} {} {}".format(
-            region, client_profile, self.headers, cache_size, cache_ttl))
+        log_msg(
+            "initialized tencent tclake client successfully {} {} {} {} {}".format(
+                region, client_profile, self.headers, cache_size, cache_ttl
+            )
+        )
 
     def _split_model_name(self, name):
         parts = name.split(".")
@@ -310,7 +314,9 @@ class TCLakeStore(AbstractStore):
             return [self.default_catalog_name, parts[0], parts[1]]
         if len(parts) == 3:
             return parts
-        raise MlflowException("invalid model name: {}, must be catalog.schema.model".format(name))
+        raise MlflowException(
+            "invalid model name: {}, must be catalog.schema.model".format(name)
+        )
 
     def _call(self, action, req):
         log_msg("req: {}\n{}".format(action, json.dumps(req, indent=2)))
@@ -336,7 +342,7 @@ class TCLakeStore(AbstractStore):
             "CatalogName": catalog_name,
             "SchemaName": schema_name,
             "ModelName": model_name,
-            "ModelVersions": self._get_model_version_numbers(name)
+            "ModelVersions": self._get_model_version_numbers(name),
         }
         resp = self._call("DescribeModelVersions", req_body)
         model_version = None
@@ -348,21 +354,32 @@ class TCLakeStore(AbstractStore):
             return None
         return _make_model_version(model_version, name)
 
-    def create_registered_model(self, name, tags=None, description=None):
+    def create_registered_model(self, name, tags=None, description=None, deployment_job_id=None):
         log_msg("create_registered_model {} {} {}".format(name, tags, description))
         [catalog_name, schema_name, model_name] = self._split_model_name(name)
         properties = []
         _add_tags_to_properties(tags, properties)
+        _add_deployment_job_id_to_properties(deployment_job_id, properties)
+        _add_workspace_id_to_properties(self.workspace_id, properties)
         req_body = {
             "CatalogName": catalog_name,
             "SchemaName": schema_name,
             "ModelName": model_name,
             "Comment": description if description else "",
-            "Properties": properties
+            "Properties": properties,
         }
-        resp = self._call("CreateModel", req_body)
+        try:
+            resp = self._call("CreateModel", req_body)
+        except Exception as e:
+            if isinstance(e, TencentCloudSDKException):
+                if e.code == "FailedOperation.MetalakeAlreadyExistsError":
+                    raise MlflowException(
+                        f"Registered Model (name={name}) already exists.",
+                        RESOURCE_ALREADY_EXISTS,
+                    )
+            raise
         return _make_model(resp["Model"])
-
+    
     def update_registered_model(self, name, description):
         log_msg("update_register_model {} {}".format(name, description))
         [catalog_name, schema_name, model_name] = self._split_model_name(name)
@@ -370,7 +387,7 @@ class TCLakeStore(AbstractStore):
             "CatalogName": catalog_name,
             "SchemaName": schema_name,
             "ModelName": model_name,
-            "NewComment": description if description else ""
+            "NewComment": description if description else "",
         }
         resp = self._call("ModifyModelComment", req_body)
         return _make_model(resp["Model"])
@@ -382,7 +399,7 @@ class TCLakeStore(AbstractStore):
             "CatalogName": catalog_name,
             "SchemaName": schema_name,
             "ModelName": model_name,
-            "NewName": new_name
+            "NewName": new_name,
         }
         resp = self._call("ModifyModelName", req_body)
         return _make_model(resp["Model"])
@@ -400,30 +417,28 @@ class TCLakeStore(AbstractStore):
             raise MlflowException("Failed to delete model {}".format(name))
 
     def _fetch_all_models(self):
-        req_body = {
-            "Offset": 0,
-            "Limit": 200,
-            "SnapshotBased": True,
-            "SnapshotId": ""
-        }
+        req_body = {"Offset": 0, "Limit": 200, "SnapshotBased": True, "SnapshotId": ""}
         resp = self._call("SearchModels", req_body)
-        total = resp['TotalCount']
-        model_list = resp['Models']
+        total = resp["TotalCount"]
+        model_list = resp["Models"]
         while len(model_list) < total:
             time.sleep(0.05)
             req_body["Offset"] += req_body["Limit"]
             req_body["SnapshotId"] = resp["SnapshotId"]
             resp = self._call("SearchModels", req_body)
-            model_list.extend(resp['Models'])
-            if len(resp['Models']) < req_body["Limit"]:
+            model_list.extend(resp["Models"])
+            if len(resp["Models"]) < req_body["Limit"]:
                 break
         return [_make_model(model) for model in model_list]
 
     def search_registered_models(
-            self, filter_string=None, max_results=None, order_by=None, page_token=None
+        self, filter_string=None, max_results=None, order_by=None, page_token=None
     ):
-        log_msg("search_registered_models {} {} {} {}".format(
-            filter_string, max_results, order_by, page_token))
+        log_msg(
+            "search_registered_models {} {} {} {}".format(
+                filter_string, max_results, order_by, page_token
+            )
+        )
         if not isinstance(max_results, int) or max_results < 1:
             raise MlflowException(
                 "Invalid value for max_results. It must be a positive integer,"
@@ -446,8 +461,16 @@ class TCLakeStore(AbstractStore):
             search_id = "model_" + str(uuid.uuid4())
             page_token = _create_page_token(0, search_id)
             self.cache[search_id] = sorted_rms
-            log_msg("search_registered_models add cache {} {} {} {} {} {}".format(
-                filter_string, max_results, order_by, page_token, search_id, len(sorted_rms)))
+            log_msg(
+                "search_registered_models add cache {} {} {} {} {} {}".format(
+                    filter_string,
+                    max_results,
+                    order_by,
+                    page_token,
+                    search_id,
+                    len(sorted_rms),
+                )
+            )
         return self._get_page_list(page_token, max_results)
 
     def get_registered_model(self, name):
@@ -472,10 +495,7 @@ class TCLakeStore(AbstractStore):
             "ModelName": model_name,
         }
         resp = self._call("DescribeModelVersions", req)
-        return [
-            _make_model_version(mv, name)
-            for mv in resp["ModelVersions"]
-        ]
+        return [_make_model_version(mv, name) for mv in resp["ModelVersions"]]
 
     def set_registered_model_tag(self, name, tag):
         log_msg("set_registered_model_tag {} {}".format(name, tag))
@@ -484,7 +504,7 @@ class TCLakeStore(AbstractStore):
             "CatalogName": catalog_name,
             "SchemaName": schema_name,
             "ModelName": model_name,
-            "Properties": _add_tag_to_properties(tag)
+            "Properties": _add_tag_to_properties(tag),
         }
         self._call("ModifyModelProperties", req)
 
@@ -495,29 +515,37 @@ class TCLakeStore(AbstractStore):
             "CatalogName": catalog_name,
             "SchemaName": schema_name,
             "ModelName": model_name,
-            "RemovedKeys": [_append_tag_key_prefix(key)]
+            "RemovedKeys": [_append_tag_key_prefix(key)],
         }
         self._call("ModifyModelProperties", req)
 
     def create_model_version(
-            self,
-            name,
-            source,
-            run_id=None,
-            tags=None,
-            run_link=None,
-            description=None,
-            local_model_path=None,
+        self,
+        name,
+        source,
+        run_id=None,
+        tags=None,
+        run_link=None,
+        description=None,
+        local_model_path=None,
+        model_id=None,
     ):
-        log_msg("create_model_version {} {} {} {} {} {} {}".format(
-            name, source, run_id, tags, run_link, description, local_model_path))
+        log_msg(
+            "create_model_version {} {} {} {} {} {} {}".format(
+                name, source, run_id, tags, run_link, description, local_model_path
+            )
+        )
         [catalog_name, schema_name, model_name] = self._split_model_name(name)
         version_alias = str(uuid.uuid4())
         properties = []
         _add_tags_to_properties(tags, properties)
         _set_run_id_to_properties(run_id, properties)
         _set_run_link_to_properties(run_link, properties)
+        _add_workspace_id_to_properties(self.workspace_id, properties)
+        _add_deployment_status_to_properties(UN_DEPLOYMENT, properties)
+        _add_model_id_to_properties(model_id, properties)
         _add_model_signature_to_properties(source, properties)
+
         req_body = {
             "CatalogName": catalog_name,
             "SchemaName": schema_name,
@@ -525,7 +553,7 @@ class TCLakeStore(AbstractStore):
             "Uri": source,
             "Comment": description if description else "",
             "Properties": properties,
-            "Aliases": [version_alias]
+            "Aliases": [version_alias],
         }
         self._call("CreateModelVersion", req_body)
         return self._get_model_version_by_alias(name, version_alias)
@@ -538,12 +566,14 @@ class TCLakeStore(AbstractStore):
             "SchemaName": schema_name,
             "ModelName": model_name,
             "ModelVersion": _set_model_version(version),
-            "NewComment": description if description else ""
+            "NewComment": description if description else "",
         }
         resp = self._call("ModifyModelVersionComment", req_body)
         return _make_model_version(resp["ModelVersion"], name)
 
-    def transition_model_version_stage(self, name, version, stage, archive_existing_versions):
+    def transition_model_version_stage(
+        self, name, version, stage, archive_existing_versions
+    ):
         raise NotImplementedError("Method not implemented")
 
     def delete_model_version(self, name, version):
@@ -557,7 +587,9 @@ class TCLakeStore(AbstractStore):
         }
         resp = self._call("DropModelVersion", req_body)
         if not resp["Dropped"]:
-            raise Exception("Failed to delete model version {} {}".format(name, version))
+            raise Exception(
+                "Failed to delete model version {} {}".format(name, version)
+            )
 
     def get_model_version(self, name, version):
         log_msg("get_model_version {} {}".format(name, version))
@@ -572,31 +604,31 @@ class TCLakeStore(AbstractStore):
         return _make_model_version(resp["ModelVersion"], name)
 
     def _fetch_all_model_versions(self):
-        req_body = {
-            "Offset": 0,
-            "Limit": 200,
-            "SnapshotBased": True,
-            "SnapshotId": ""
-        }
+        req_body = {"Offset": 0, "Limit": 200, "SnapshotBased": True, "SnapshotId": ""}
         resp = self._call("SearchModelVersions", req_body)
-        total = resp['TotalCount']
-        model_version_list = resp['ModelVersions']
+        total = resp["TotalCount"]
+        model_version_list = resp["ModelVersions"]
         while len(model_version_list) < total:
             time.sleep(0.05)
             req_body["Offset"] += req_body["Limit"]
             req_body["SnapshotId"] = resp["SnapshotId"]
             resp = self._call("SearchModelVersions", req_body)
-            model_version_list.extend(resp['ModelVersions'])
-            if len(resp['ModelVersions']) < req_body["Limit"]:
+            model_version_list.extend(resp["ModelVersions"])
+            if len(resp["ModelVersions"]) < req_body["Limit"]:
                 break
-        return [_make_model_version(model_version, _get_model_version_name(model_version)) for model_version in
-                model_version_list]
+        return [
+            _make_model_version(model_version, _get_model_version_name(model_version))
+            for model_version in model_version_list
+        ]
 
     def search_model_versions(
-            self, filter_string=None, max_results=None, order_by=None, page_token=None
+        self, filter_string=None, max_results=None, order_by=None, page_token=None
     ):
-        log_msg("search_model_versions {} {} {} {}".format(
-            filter_string, max_results, order_by, page_token))
+        log_msg(
+            "search_model_versions {} {} {} {}".format(
+                filter_string, max_results, order_by, page_token
+            )
+        )
         if not isinstance(max_results, int) or max_results < 1:
             raise MlflowException(
                 "Invalid value for max_results. It must be a positive integer,"
@@ -615,39 +647,59 @@ class TCLakeStore(AbstractStore):
             filtered_mvs = SearchModelVersionUtils.filter(model_versions, filter_string)
             sorted_mvs = SearchModelVersionUtils.sort(
                 filtered_mvs,
-                order_by or ["last_updated_timestamp DESC", "name ASC", "version_number DESC"],
+                order_by
+                or ["last_updated_timestamp DESC", "name ASC", "version_number DESC"],
             )
             if len(sorted_mvs) == 0:
                 return PagedList([], None)
             search_id = "model_version_" + str(uuid.uuid4())
             page_token = _create_page_token(0, search_id)
             self.cache[search_id] = sorted_mvs
-            log_msg("search_model_versions add cache {} {} {} {} {} {}".format(
-                filter_string, max_results, order_by, page_token, search_id, len(sorted_mvs)))
+            log_msg(
+                "search_model_versions add cache {} {} {} {} {} {}".format(
+                    filter_string,
+                    max_results,
+                    order_by,
+                    page_token,
+                    search_id,
+                    len(sorted_mvs),
+                )
+            )
 
         return self._get_page_list(page_token, max_results)
 
     def _get_page_list(self, page_token, max_results):
         token_info = _parse_page_token(page_token)
-        log_msg("_get_page_list token_info {} {} {}".format(
-            page_token, max_results, token_info))
-        sorted_mvs = self.cache.get(token_info['search_id'])
+        log_msg(
+            "_get_page_list token_info {} {} {}".format(
+                page_token, max_results, token_info
+            )
+        )
+        sorted_mvs = self.cache.get(token_info["search_id"])
         if sorted_mvs is None:
             raise MlflowException(
                 "Invalid page token: search id not found or expired",
                 INVALID_PARAMETER_VALUE,
             )
-        start_offset = token_info['offset']
+        start_offset = token_info["offset"]
         final_offset = start_offset + max_results
 
-        paginated_rms = sorted_mvs[start_offset: min(len(sorted_mvs), final_offset)]
+        paginated_rms = sorted_mvs[start_offset : min(len(sorted_mvs), final_offset)]
         next_page_token = None
         if final_offset < len(sorted_mvs):
-            next_page_token = _create_page_token(final_offset, token_info['search_id'])
+            next_page_token = _create_page_token(final_offset, token_info["search_id"])
         else:
-            self.cache.pop(token_info['search_id'], None)
-            log_msg("pop cache {} {} {} {} {} {}".format(
-                token_info['search_id'], start_offset, final_offset, len(sorted_mvs), page_token, next_page_token))
+            self.cache.pop(token_info["search_id"], None)
+            log_msg(
+                "pop cache {} {} {} {} {} {}".format(
+                    token_info["search_id"],
+                    start_offset,
+                    final_offset,
+                    len(sorted_mvs),
+                    page_token,
+                    next_page_token,
+                )
+            )
         return PagedList(paginated_rms, next_page_token)
 
     def set_model_version_tag(self, name, version, tag):
@@ -658,7 +710,7 @@ class TCLakeStore(AbstractStore):
             "SchemaName": schema_name,
             "ModelName": model_name,
             "ModelVersion": _set_model_version(version),
-            "Properties": _add_tag_to_properties(tag)
+            "Properties": _add_tag_to_properties(tag),
         }
         self._call("ModifyModelVersionProperties", req)
 
@@ -670,7 +722,7 @@ class TCLakeStore(AbstractStore):
             "SchemaName": schema_name,
             "ModelName": model_name,
             "ModelVersion": _set_model_version(version),
-            "RemovedKeys": [_append_tag_key_prefix(key)]
+            "RemovedKeys": [_append_tag_key_prefix(key)],
         }
         self._call("ModifyModelVersionProperties", req)
 
@@ -687,4 +739,5 @@ class TCLakeStore(AbstractStore):
     def get_model_version_download_uri(self, name, version):
         log_msg("get_model_version_download_uri {} {}".format(name, version))
         model_version = self.get_model_version(name, version)
-        return model_version.source
+        return model_version.source
+

{mlflow_tclake_plugin-2.0.1.dist-info → mlflow_tclake_plugin-3.0.1.dist-info}/METADATA

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: mlflow-tclake-plugin
-Version: 2.0.1
+Version: 3.0.1
 Summary: Tclake plugin for MLflow
 License-File: LICENSE.txt
-Requires-Dist: mlflow>=2.7.2
+Requires-Dist: mlflow>=3.1.4
 Requires-Dist: tencentcloud-sdk-python-common>=3.0.1478
 Dynamic: license-file
 Dynamic: requires-dist

mlflow_tclake_plugin-3.0.1.dist-info/RECORD

@@ -0,0 +1,8 @@
+mlflow_tclake_plugin/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mlflow_tclake_plugin/tclake_store.py,sha256=Ys3giqvpdth7gwZ7O3TeaQH1OJH08XUpqjRkOVIM9JM,28301
+mlflow_tclake_plugin-3.0.1.dist-info/licenses/LICENSE.txt,sha256=X60Z7_gpe--AyXGUWeFWpOpFl5m-yeemfodWGJn1rUA,1067
+mlflow_tclake_plugin-3.0.1.dist-info/METADATA,sha256=tz2ZxO4mTJTZR6WlANcobb91uTAvyWCS6Ff0RDuR1Cs,271
+mlflow_tclake_plugin-3.0.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+mlflow_tclake_plugin-3.0.1.dist-info/entry_points.txt,sha256=xKXky9-NyJWEG1SgBknrqskN2rsZG3t4UXiSzkTcsuE,85
+mlflow_tclake_plugin-3.0.1.dist-info/top_level.txt,sha256=zrA5UNyfF3skRmxPNsvrJI3yf-um6CJX_xO5KvWc2o0,21
+mlflow_tclake_plugin-3.0.1.dist-info/RECORD,,

mlflow_tclake_plugin-2.0.1.dist-info/RECORD

@@ -1,8 +0,0 @@
-mlflow_tclake_plugin/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mlflow_tclake_plugin/tclake_store.py,sha256=djAE2-C_2BKkguy8WN-uETIOcC9qLiiCSDu7teXFRQc,27054
-mlflow_tclake_plugin-2.0.1.dist-info/licenses/LICENSE.txt,sha256=X60Z7_gpe--AyXGUWeFWpOpFl5m-yeemfodWGJn1rUA,1067
-mlflow_tclake_plugin-2.0.1.dist-info/METADATA,sha256=rVelcV3Jsjdi0RVsbdGPkYGONGiONW63dL6iJb1rzFc,271
-mlflow_tclake_plugin-2.0.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-mlflow_tclake_plugin-2.0.1.dist-info/entry_points.txt,sha256=xKXky9-NyJWEG1SgBknrqskN2rsZG3t4UXiSzkTcsuE,85
-mlflow_tclake_plugin-2.0.1.dist-info/top_level.txt,sha256=zrA5UNyfF3skRmxPNsvrJI3yf-um6CJX_xO5KvWc2o0,21
-mlflow_tclake_plugin-2.0.1.dist-info/RECORD,,