mlflow-tclake-plugin 2.1.7.dev2__tar.gz → 2.1.7.dev4__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: mlflow-tclake-plugin
3
- Version: 2.1.7.dev2
3
+ Version: 2.1.7.dev4
4
4
  Summary: Tclake plugin for MLflow
5
5
  License-File: LICENSE.txt
6
6
  Requires-Dist: mlflow==2.17.2
@@ -0,0 +1,466 @@
1
+ """
2
+ 腾讯云凭证提供器(自包含,不依赖 wedata_automl)
3
+
4
+ 密钥优先级(高 -> 低):
5
+
6
+ 1. 固定密钥:`TENCENTCLOUD_SECRET_ID` + `TENCENTCLOUD_SECRET_KEY` 且**没有**
7
+ `TENCENTCLOUD_TOKEN` —— 视为长期密钥,直接使用,不做续期、不启动定时任务。
8
+ 2. 临时密钥:以环境变量中的密钥(含 `TENCENTCLOUD_TOKEN`)作为引导凭证,
9
+ 用 `CommonClient` 调用 wedata 云 API `DescribeCamTemporaryToken` 换取临时密钥,
10
+ 并由后台定时任务在**过期前 60 秒**自动续期。
11
+
12
+ 为什么续期必须用定时任务,而不是"用的时候再检查"
13
+ --------------------------------------------------
14
+ 续期本身也是一次需要签名的云 API 调用,用的是当前这把临时密钥。一旦它彻底过期,
15
+ 续期请求会因鉴权失败而无法完成,凭证链就永久断裂。所以必须在过期前主动续期 ——
16
+ 即使这段时间内业务侧完全没有调用云 API。
17
+
18
+ ProductName 说明
19
+ ----------------
20
+ `DescribeCamTemporaryToken` 的 `ProductName` 仅支持 `wedata` 与 `tccatalog`,
21
+ 服务端据此下发 `name/{ProductName}:*` 的授权策略,两者的 token 不能互用。
22
+ 本插件访问 tccatalog,因此固定使用 `tccatalog`。
23
+
24
+ 注意:`DescribeCamTemporaryToken` 本身始终发往 **wedata** 服务端点
25
+ (`TENCENTCLOUD_ENDPOINT` 指向的是 tccatalog,不能用于这个调用)。
26
+
27
+ 注意:本模块刻意不依赖 `wedata_automl`——本插件是被 wedata_automl 安装使用的,
28
+ 反向依赖会形成循环。
29
+ """
30
+
31
+ import os
32
+ import threading
33
+ import time
34
+
35
+ ENV_SECRET_ID = "TENCENTCLOUD_SECRET_ID"
36
+ ENV_SECRET_KEY = "TENCENTCLOUD_SECRET_KEY"
37
+ ENV_TOKEN = "TENCENTCLOUD_TOKEN"
38
+ ENV_TMP_TOKEN = "TENCENTCLOUD_TMP_TOKEN"
39
+ ENV_PROJECT_ID = "WEDATA_PROJECT_ID"
40
+ ENV_WEDATA_ENDPOINT = "TENCENT_CLOUD_SDK_ENDPOINT"
41
+ ENV_REFRESH_AHEAD = "WEDATA_SECRET_REFRESH_AHEAD_SECONDS"
42
+ ENV_DISABLE_DYNAMIC = "WEDATA_DISABLE_DYNAMIC_SECRET"
43
+
44
+ PRODUCT_WEDATA = "wedata"
45
+ PRODUCT_TCCATALOG = "tccatalog"
46
+
47
+ WEDATA_SERVICE = "wedata"
48
+ WEDATA_API_VERSION = "2021-08-20"
49
+ DEFAULT_WEDATA_ENDPOINT = "wedata.internal.tencentcloudapi.com"
50
+
51
+ DEFAULT_REFRESH_AHEAD_SECONDS = 60
52
+ MIN_REFRESH_DELAY_SECONDS = 5
53
+ RETRY_DELAY_SECONDS = 30
54
+ # ExpiredTime 合法区间(秒级时间戳):2020-01-01 ~ 2100-01-01
55
+ MIN_VALID_EPOCH_SECONDS = 1577836800
56
+ MAX_VALID_EPOCH_SECONDS = 4102444800
57
+
58
+
59
+ def _debug(msg):
60
+ if os.getenv("TENCENTCLOUD_DEBUG", None):
61
+ print(msg)
62
+
63
+
64
+ def mask_secret_id(secret_id):
65
+ """脱敏展示 SecretId,仅保留前后各 4 位"""
66
+ if not secret_id:
67
+ return ""
68
+ if len(secret_id) <= 8:
69
+ return "*" * len(secret_id)
70
+ return "{}{}{}".format(secret_id[:4], "*" * 6, secret_id[-4:])
71
+
72
+
73
+ def _read_refresh_ahead():
74
+ try:
75
+ value = int(os.getenv(ENV_REFRESH_AHEAD, str(DEFAULT_REFRESH_AHEAD_SECONDS)))
76
+ return value if value >= 0 else DEFAULT_REFRESH_AHEAD_SECONDS
77
+ except (TypeError, ValueError):
78
+ return DEFAULT_REFRESH_AHEAD_SECONDS
79
+
80
+
81
+ def _dynamic_disabled():
82
+ """是否禁用"调云 API 换临时密钥"(逃生阀)"""
83
+ return str(os.getenv(ENV_DISABLE_DYNAMIC, "")).lower() in ("1", "true", "yes")
84
+
85
+
86
+ def _parse_expire(value):
87
+ """
88
+ 解析 ExpiredTime
89
+
90
+ `DescribeCamTemporaryToken` 返回的 `ExpiredTime` 是**秒级绝对时间戳**
91
+ (例如 1788789863 -> 2026-09-07 22:04:23)。
92
+
93
+ 只做解析与合法性校验,不做语义猜测:异常值告警并返回 0,
94
+ 由上层按"没有到期信息"处理,避免静默换算把问题藏起来。
95
+ """
96
+ if value is None:
97
+ return 0
98
+ try:
99
+ expire = int(value)
100
+ except (TypeError, ValueError):
101
+ _debug("invalid ExpiredTime (not an int): {!r}".format(value))
102
+ return 0
103
+ if expire <= 0:
104
+ _debug("invalid ExpiredTime (non-positive): {}".format(expire))
105
+ return 0
106
+ if expire < MIN_VALID_EPOCH_SECONDS or expire > MAX_VALID_EPOCH_SECONDS:
107
+ _debug("ExpiredTime {} is not a plausible epoch-seconds value; "
108
+ "treating as unknown".format(expire))
109
+ return 0
110
+ return expire
111
+
112
+
113
+ class TempCredential(object):
114
+ """密钥快照"""
115
+
116
+ __slots__ = ("secret_id", "secret_key", "token", "expire", "owner_uin", "user_id")
117
+
118
+ def __init__(self, secret_id=None, secret_key=None, token=None, expire=0,
119
+ owner_uin=None, user_id=None):
120
+ self.secret_id = secret_id or None
121
+ self.secret_key = secret_key or None
122
+ self.token = token or None
123
+ self.expire = int(expire or 0)
124
+ self.owner_uin = owner_uin or None
125
+ self.user_id = user_id or None
126
+
127
+ @property
128
+ def valid(self):
129
+ return bool(self.secret_id and self.secret_key)
130
+
131
+ @property
132
+ def is_static(self):
133
+ """无 token 即视为固定(长期)密钥"""
134
+ return self.valid and not self.token
135
+
136
+ def identity(self):
137
+ return self.secret_id, self.secret_key, self.token, self.expire
138
+
139
+ def is_expiring(self, refresh_ahead_seconds=DEFAULT_REFRESH_AHEAD_SECONDS):
140
+ if self.expire <= 0:
141
+ return False
142
+ return time.time() >= (self.expire - max(refresh_ahead_seconds, 0))
143
+
144
+ def seconds_to_expire(self):
145
+ if self.expire <= 0:
146
+ return None
147
+ return int(self.expire - time.time())
148
+
149
+ def __repr__(self): # 不输出 SecretKey / Token
150
+ return "TempCredential(secret_id={}, has_token={}, expire={}, ttl={})".format(
151
+ mask_secret_id(self.secret_id), bool(self.token), self.expire,
152
+ self.seconds_to_expire()
153
+ )
154
+
155
+
156
+ class CredentialProvider(object):
157
+ """
158
+ 凭证提供器
159
+
160
+ - 固定密钥:直接返回,不启动定时任务
161
+ - 临时密钥:调用 `DescribeCamTemporaryToken` 换取,并由后台定时任务提前续期
162
+ """
163
+
164
+ def __init__(self, product_name=PRODUCT_TCCATALOG, project_id=None, region=None,
165
+ refresh_ahead_seconds=None):
166
+ self._product_name = product_name or PRODUCT_TCCATALOG
167
+ self._project_id = project_id or os.getenv(ENV_PROJECT_ID) or None
168
+ self._region = region
169
+ self._refresh_ahead_seconds = (
170
+ refresh_ahead_seconds if refresh_ahead_seconds is not None else _read_refresh_ahead()
171
+ )
172
+ self._base = TempCredential(
173
+ secret_id=os.getenv(ENV_SECRET_ID) or None,
174
+ secret_key=os.getenv(ENV_SECRET_KEY) or None,
175
+ token=os.getenv(ENV_TOKEN) or os.getenv(ENV_TMP_TOKEN) or None,
176
+ expire=0,
177
+ )
178
+ self._cached = None
179
+ self._timer = None
180
+ self._stopped = False
181
+ self._lock = threading.RLock()
182
+
183
+ @property
184
+ def product_name(self):
185
+ return self._product_name
186
+
187
+ @property
188
+ def use_static_secret(self):
189
+ """固定密钥(无 token),或逃生阀开启 -> 不调云 API 换密钥"""
190
+ return self._base.is_static or _dynamic_disabled()
191
+
192
+ def get(self, force_refresh=False):
193
+ """获取当前可用密钥;临时密钥在过期/即将过期时同步续期一次(兜底)"""
194
+ if self.use_static_secret:
195
+ return self._base
196
+
197
+ with self._lock:
198
+ cached = self._cached
199
+ if (
200
+ not force_refresh
201
+ and cached is not None
202
+ and cached.valid
203
+ and not cached.is_expiring(self._refresh_ahead_seconds)
204
+ ):
205
+ return cached
206
+
207
+ refreshed = self._refresh_locked()
208
+ if refreshed is not None:
209
+ return refreshed
210
+ if cached is not None and cached.valid:
211
+ return cached
212
+ return self._base
213
+
214
+ def start_auto_refresh(self):
215
+ """启动后台定时续期任务(幂等);固定密钥场景不创建线程"""
216
+ if self.use_static_secret:
217
+ return False
218
+
219
+ with self._lock:
220
+ self._stopped = False
221
+ if self._timer is not None:
222
+ return True
223
+ self._refresh_locked()
224
+ self._schedule_locked()
225
+ return True
226
+
227
+ def stop_auto_refresh(self):
228
+ with self._lock:
229
+ self._stopped = True
230
+ if self._timer is not None:
231
+ self._timer.cancel()
232
+ self._timer = None
233
+
234
+ def _schedule_locked(self):
235
+ if self._stopped:
236
+ return
237
+ cached = self._cached
238
+ if cached is not None and cached.valid and cached.expire > 0:
239
+ ttl = cached.expire - time.time()
240
+ delay = ttl - self._refresh_ahead_seconds
241
+ if delay < MIN_REFRESH_DELAY_SECONDS:
242
+ # TTL 比提前量还短:固定用最小延迟会变成每 5s 打一次云 API(可能限频),
243
+ # 退化为按 TTL 的一半续期
244
+ delay = max(ttl / 2.0, MIN_REFRESH_DELAY_SECONDS)
245
+ _debug("credential ttl ({}s) shorter than refresh-ahead ({}s), "
246
+ "next refresh in {}s".format(
247
+ int(ttl), self._refresh_ahead_seconds, int(delay)))
248
+ else:
249
+ delay = RETRY_DELAY_SECONDS
250
+
251
+ timer = threading.Timer(delay, self._on_timer)
252
+ timer.daemon = True
253
+ timer.name = "tclake-cred-refresh"
254
+ timer.start()
255
+ self._timer = timer
256
+ _debug("next credential refresh in {}s".format(int(delay)))
257
+
258
+ def _on_timer(self):
259
+ with self._lock:
260
+ self._timer = None
261
+ if self._stopped:
262
+ return
263
+ try:
264
+ self._refresh_locked()
265
+ except Exception as e:
266
+ _debug("scheduled credential refresh failed: {}".format(e))
267
+ self._schedule_locked()
268
+
269
+ def _refresh_locked(self):
270
+ fetched = self._fetch_cam_token()
271
+ if fetched is None or not fetched.valid:
272
+ return None
273
+ changed = self._cached is None or self._cached.identity() != fetched.identity()
274
+ self._cached = fetched
275
+ self.sync_tencentcloud_env(fetched)
276
+ if changed:
277
+ _debug("cam temporary credential refreshed: {}".format(fetched))
278
+ return fetched
279
+
280
+ def _signing_credential(self):
281
+ """构造调用 DescribeCamTemporaryToken 的签名凭证(链式续期)"""
282
+ from tencentcloud.common import credential as tc_credential
283
+
284
+ cached = self._cached
285
+ if cached is not None and cached.valid and not cached.is_expiring(0):
286
+ source = cached
287
+ else:
288
+ source = self._base
289
+ return tc_credential.Credential(source.secret_id, source.secret_key, source.token)
290
+
291
+ def _fetch_cam_token(self):
292
+ if not self._base.valid:
293
+ _debug("cannot fetch cam temporary token: secret id/key not set")
294
+ return None
295
+ if not self._project_id:
296
+ _debug("cannot fetch cam temporary token: WEDATA_PROJECT_ID not set")
297
+ return None
298
+ if not self._region:
299
+ _debug("cannot fetch cam temporary token: region unknown")
300
+ return None
301
+
302
+ try:
303
+ from tencentcloud.common.common_client import CommonClient
304
+
305
+ client = CommonClient(
306
+ WEDATA_SERVICE,
307
+ WEDATA_API_VERSION,
308
+ self._signing_credential(),
309
+ self._region,
310
+ self._build_wedata_profile(),
311
+ )
312
+ params = {"ProjectId": self._project_id, "ProductName": self._product_name}
313
+ resp = client.call_json("DescribeCamTemporaryToken", params)
314
+ except Exception as e:
315
+ _debug("call DescribeCamTemporaryToken failed: {}".format(e))
316
+ return None
317
+
318
+ data = (resp or {}).get("Response", {}).get("Data") or {}
319
+ secret_id = data.get("SecretId")
320
+ secret_key = data.get("SecretKey")
321
+ if not secret_id or not secret_key:
322
+ _debug("DescribeCamTemporaryToken returned no credential")
323
+ return None
324
+
325
+ return TempCredential(
326
+ secret_id=secret_id,
327
+ secret_key=secret_key,
328
+ token=data.get("Token"),
329
+ expire=_normalize_expire(data.get("ExpiredTime")),
330
+ owner_uin=data.get("OwnerUserId"),
331
+ user_id=data.get("UserId"),
332
+ )
333
+
334
+ @staticmethod
335
+ def _build_wedata_profile():
336
+ """
337
+ 构造调用 wedata 云 API 的 ClientProfile
338
+
339
+ 不能复用 TENCENTCLOUD_ENDPOINT——那个指向 tccatalog。
340
+ """
341
+ from tencentcloud.common.profile.client_profile import ClientProfile
342
+ from tencentcloud.common.profile.http_profile import HttpProfile
343
+
344
+ http_profile = HttpProfile()
345
+ http_profile.protocol = "https"
346
+ http_profile.endpoint = os.getenv(ENV_WEDATA_ENDPOINT) or DEFAULT_WEDATA_ENDPOINT
347
+ client_profile = ClientProfile()
348
+ client_profile.httpProfile = http_profile
349
+ return client_profile
350
+
351
+ @staticmethod
352
+ def sync_tencentcloud_env(cred):
353
+ """把最新密钥同步到 TENCENTCLOUD_* 环境变量"""
354
+ if not cred.valid:
355
+ return
356
+ os.environ[ENV_SECRET_ID] = cred.secret_id
357
+ os.environ[ENV_SECRET_KEY] = cred.secret_key
358
+ if cred.token:
359
+ os.environ[ENV_TOKEN] = cred.token
360
+ os.environ[ENV_TMP_TOKEN] = cred.token
361
+
362
+ def as_sdk_credential(self):
363
+ """包装为腾讯云 SDK 可直接使用的动态凭证"""
364
+ return DynamicCredential(self)
365
+
366
+
367
+ class DynamicCredential(object):
368
+ """
369
+ 腾讯云 SDK 动态凭证
370
+
371
+ SDK 每次请求签名(含每次重试)都会调用 `get_credential_info()`,
372
+ 因此 client 无需重建即可用上续期后的密钥。
373
+ """
374
+
375
+ def __init__(self, provider):
376
+ self._provider = provider
377
+
378
+ def _current(self):
379
+ cred = self._provider.get()
380
+ if not cred.valid:
381
+ raise _credential_error(
382
+ "no valid tencentcloud credential: {}/{} not set and "
383
+ "DescribeCamTemporaryToken unavailable".format(ENV_SECRET_ID, ENV_SECRET_KEY)
384
+ )
385
+ return cred
386
+
387
+ def get_credential_info(self):
388
+ cred = self._current()
389
+ return cred.secret_id, cred.secret_key, cred.token
390
+
391
+ @property
392
+ def secret_id(self):
393
+ return self._current().secret_id
394
+
395
+ @property
396
+ def secretId(self):
397
+ return self.secret_id
398
+
399
+ @property
400
+ def secret_key(self):
401
+ return self._current().secret_key
402
+
403
+ @property
404
+ def secretKey(self):
405
+ return self.secret_key
406
+
407
+ @property
408
+ def token(self):
409
+ return self._current().token
410
+
411
+ @property
412
+ def provider(self):
413
+ return self._provider
414
+
415
+ def __repr__(self): # 不输出 SecretKey / Token
416
+ return "DynamicCredential(product={}, static={})".format(
417
+ self._provider.product_name, self._provider.use_static_secret
418
+ )
419
+
420
+
421
+ def _credential_error(message):
422
+ try:
423
+ from mlflow.exceptions import MlflowException
424
+
425
+ return MlflowException(message)
426
+ except ImportError:
427
+ return ValueError(message)
428
+
429
+
430
+ _providers = {}
431
+ _providers_lock = threading.Lock()
432
+
433
+
434
+ def get_default_provider(product_name=PRODUCT_TCCATALOG, project_id=None, region=None,
435
+ auto_refresh=True):
436
+ """
437
+ 获取(并按需创建)进程级共享的 Provider
438
+
439
+ 按 (product_name, project_id, region) 维度缓存:wedata 与 tccatalog 的
440
+ token 授权范围不同,不能共用。
441
+ """
442
+ key = (
443
+ product_name or PRODUCT_TCCATALOG,
444
+ project_id or os.getenv(ENV_PROJECT_ID) or None,
445
+ region,
446
+ )
447
+ with _providers_lock:
448
+ provider = _providers.get(key)
449
+ if provider is None:
450
+ provider = CredentialProvider(
451
+ product_name=key[0], project_id=key[1], region=key[2]
452
+ )
453
+ _providers[key] = provider
454
+ if auto_refresh:
455
+ provider.start_auto_refresh()
456
+ return provider
457
+
458
+
459
+ def reset_providers():
460
+ """停止并清空所有 Provider(主要用于测试)"""
461
+ global _providers
462
+ with _providers_lock:
463
+ providers = list(_providers.values())
464
+ _providers = {}
465
+ for provider in providers:
466
+ provider.stop_auto_refresh()
@@ -6,10 +6,11 @@ import uuid
6
6
  from datetime import datetime
7
7
 
8
8
  from cachetools import TTLCache
9
- from tencentcloud.common import credential
10
9
  from tencentcloud.common.common_client import CommonClient
11
10
  from tencentcloud.common.profile.client_profile import ClientProfile
12
11
 
12
+ from mlflow_tclake_plugin.credential_provider import PRODUCT_TCCATALOG, get_default_provider
13
+
13
14
  from mlflow.entities.model_registry import RegisteredModel, ModelVersion, ModelVersionTag
14
15
  from mlflow.exceptions import MlflowException
15
16
  from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE, RESOURCE_ALREADY_EXISTS
@@ -302,19 +303,24 @@ class TCLakeStore(AbstractStore):
302
303
  def __init__(self, store_uri=None, tracking_uri=None):
303
304
  super().__init__(store_uri, tracking_uri)
304
305
  log_msg("initializing tencent tclake client {} {}".format(store_uri, tracking_uri))
305
- sid = os.getenv("TENCENTCLOUD_SECRET_ID", "")
306
- if len(sid) == 0:
307
- raise MlflowException("TENCENTCLOUD_SECRET_ID is not set")
308
- sk = os.getenv("TENCENTCLOUD_SECRET_KEY", "")
309
- if len(sk) == 0:
310
- raise MlflowException("TENCENTCLOUD_SECRET_KEY is not set")
311
- token = os.getenv("TENCENTCLOUD_TOKEN", None)
312
- client_profile = get_tencent_cloud_client_profile()
313
- cred = credential.Credential(sid, sk, token)
314
- parts = store_uri.split(":")
306
+ parts = (store_uri or "").split(":")
315
307
  if len(parts) < 2:
316
308
  raise MlflowException("set store_uri tclake:{region}")
317
309
  region = parts[1]
310
+ # 凭证优先级:固定密钥 > DescribeCamTemporaryToken 换取的临时密钥。
311
+ # 初始化即启动定时续期(过期前 60s),并包装成动态凭证——
312
+ # SDK 每次请求签名都会重新取值,client 无需重建。
313
+ self.credential_provider = get_default_provider(
314
+ product_name=PRODUCT_TCCATALOG, region=region
315
+ )
316
+ self.credential_provider.start_auto_refresh()
317
+ cloud_credential = self.credential_provider.get()
318
+ if not cloud_credential.secret_id:
319
+ raise MlflowException("TENCENTCLOUD_SECRET_ID is not set")
320
+ if not cloud_credential.secret_key:
321
+ raise MlflowException("TENCENTCLOUD_SECRET_KEY is not set")
322
+ client_profile = get_tencent_cloud_client_profile()
323
+ cred = self.credential_provider.as_sdk_credential()
318
324
  self.client = CommonClient("tccatalog", "2024-10-24", cred, region, client_profile)
319
325
  self.headers = get_tencent_cloud_headers()
320
326
  self.default_catalog_name = os.getenv("TENCENTCLOUD_DEFAULT_CATALOG_NAME", "default")
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: mlflow-tclake-plugin
3
- Version: 2.1.7.dev2
3
+ Version: 2.1.7.dev4
4
4
  Summary: Tclake plugin for MLflow
5
5
  License-File: LICENSE.txt
6
6
  Requires-Dist: mlflow==2.17.2
@@ -2,6 +2,7 @@ LICENSE.txt
2
2
  README.md
3
3
  setup.py
4
4
  mlflow_tclake_plugin/__init__.py
5
+ mlflow_tclake_plugin/credential_provider.py
5
6
  mlflow_tclake_plugin/tclake_store.py
6
7
  mlflow_tclake_plugin.egg-info/PKG-INFO
7
8
  mlflow_tclake_plugin.egg-info/SOURCES.txt
@@ -2,7 +2,7 @@ from setuptools import find_packages, setup
2
2
 
3
3
  setup(
4
4
  name="mlflow-tclake-plugin",
5
- version="2.1.7.dev2",
5
+ version="2.1.7.dev4",
6
6
  description="Tclake plugin for MLflow",
7
7
  packages=find_packages(),
8
8
  # Require MLflow as a dependency of the plugin, so that plugin users can simply install