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

mlflow_tclake_plugin/tclake_store.py

@@ -3,7 +3,6 @@ import json
 import os
 import time
 import uuid
-import urllib.parse
 from datetime import datetime
 
 from cachetools import TTLCache
@@ -11,23 +10,16 @@ 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, RESOURCE_ALREADY_EXISTS
-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
+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:
@@ -73,13 +65,9 @@ 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_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_UUID_KEY = "tccatalog.identifier"
 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:
@@ -98,6 +86,10 @@ 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)
 
@@ -113,15 +105,10 @@ 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
 
 
@@ -131,47 +118,84 @@ 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 _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 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_signature_to_properties(source, properties):
     log_msg("model source is {}".format(source))
-    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))
+    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([])
     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
 
 
@@ -195,20 +219,17 @@ 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"]),
@@ -219,8 +240,6 @@ 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",
     )
 
@@ -248,9 +267,7 @@ 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
@@ -261,50 +278,29 @@ 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(".")
@@ -314,9 +310,7 @@ 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)))
@@ -342,7 +336,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
@@ -354,32 +348,21 @@ class TCLakeStore(AbstractStore):
             return None
         return _make_model_version(model_version, name)
 
-    def create_registered_model(self, name, tags=None, description=None, deployment_job_id=None):
+    def create_registered_model(self, name, tags=None, description=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
         }
-        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
+        resp = self._call("CreateModel", req_body)
         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)
@@ -387,7 +370,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"])
@@ -399,7 +382,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"])
@@ -417,28 +400,30 @@ 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,"
@@ -461,16 +446,8 @@ 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):
@@ -495,7 +472,10 @@ 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))
@@ -504,7 +484,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)
 
@@ -515,37 +495,29 @@ 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,
-        model_id=None,
+            self,
+            name,
+            source,
+            run_id=None,
+            tags=None,
+            run_link=None,
+            description=None,
+            local_model_path=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,
@@ -553,7 +525,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)
@@ -566,14 +538,12 @@ 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):
@@ -587,9 +557,7 @@ 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))
@@ -604,31 +572,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,"
@@ -647,59 +615,39 @@ 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):
@@ -710,7 +658,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)
 
@@ -722,7 +670,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)
 
@@ -739,5 +687,4 @@ 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-0.0.1.dist-info → mlflow_tclake_plugin-2.0.1.dist-info}/METADATA

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

mlflow_tclake_plugin-2.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=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,,

mlflow_tclake_plugin-2.0.1.dist-info/licenses/LICENSE.txt

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) [2025] [tencent]
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

mlflow_tclake_plugin-0.0.1.dist-info/RECORD

@@ -1,7 +0,0 @@
-mlflow_tclake_plugin/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mlflow_tclake_plugin/tclake_store.py,sha256=Ys3giqvpdth7gwZ7O3TeaQH1OJH08XUpqjRkOVIM9JM,28301
-mlflow_tclake_plugin-0.0.1.dist-info/METADATA,sha256=My3rL7JEx5v_HsA42LimQvwHJU1ms0Qa2mvHlYGZ70s,223
-mlflow_tclake_plugin-0.0.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-mlflow_tclake_plugin-0.0.1.dist-info/entry_points.txt,sha256=xKXky9-NyJWEG1SgBknrqskN2rsZG3t4UXiSzkTcsuE,85
-mlflow_tclake_plugin-0.0.1.dist-info/top_level.txt,sha256=zrA5UNyfF3skRmxPNsvrJI3yf-um6CJX_xO5KvWc2o0,21
-mlflow_tclake_plugin-0.0.1.dist-info/RECORD,,