genesis-flow 1.0.0__py3-none-any.whl

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.
Files changed (645) hide show
  1. genesis_flow-1.0.0.dist-info/METADATA +822 -0
  2. genesis_flow-1.0.0.dist-info/RECORD +645 -0
  3. genesis_flow-1.0.0.dist-info/WHEEL +5 -0
  4. genesis_flow-1.0.0.dist-info/entry_points.txt +19 -0
  5. genesis_flow-1.0.0.dist-info/licenses/LICENSE.txt +202 -0
  6. genesis_flow-1.0.0.dist-info/top_level.txt +1 -0
  7. mlflow/__init__.py +367 -0
  8. mlflow/__main__.py +3 -0
  9. mlflow/ag2/__init__.py +56 -0
  10. mlflow/ag2/ag2_logger.py +294 -0
  11. mlflow/anthropic/__init__.py +40 -0
  12. mlflow/anthropic/autolog.py +129 -0
  13. mlflow/anthropic/chat.py +144 -0
  14. mlflow/artifacts/__init__.py +268 -0
  15. mlflow/autogen/__init__.py +144 -0
  16. mlflow/autogen/chat.py +142 -0
  17. mlflow/azure/__init__.py +26 -0
  18. mlflow/azure/auth_handler.py +257 -0
  19. mlflow/azure/client.py +319 -0
  20. mlflow/azure/config.py +120 -0
  21. mlflow/azure/connection_factory.py +340 -0
  22. mlflow/azure/exceptions.py +27 -0
  23. mlflow/azure/stores.py +327 -0
  24. mlflow/azure/utils.py +183 -0
  25. mlflow/bedrock/__init__.py +45 -0
  26. mlflow/bedrock/_autolog.py +202 -0
  27. mlflow/bedrock/chat.py +122 -0
  28. mlflow/bedrock/stream.py +160 -0
  29. mlflow/bedrock/utils.py +43 -0
  30. mlflow/cli.py +707 -0
  31. mlflow/client.py +12 -0
  32. mlflow/config/__init__.py +56 -0
  33. mlflow/crewai/__init__.py +79 -0
  34. mlflow/crewai/autolog.py +253 -0
  35. mlflow/crewai/chat.py +29 -0
  36. mlflow/data/__init__.py +75 -0
  37. mlflow/data/artifact_dataset_sources.py +170 -0
  38. mlflow/data/code_dataset_source.py +40 -0
  39. mlflow/data/dataset.py +123 -0
  40. mlflow/data/dataset_registry.py +168 -0
  41. mlflow/data/dataset_source.py +110 -0
  42. mlflow/data/dataset_source_registry.py +219 -0
  43. mlflow/data/delta_dataset_source.py +167 -0
  44. mlflow/data/digest_utils.py +108 -0
  45. mlflow/data/evaluation_dataset.py +562 -0
  46. mlflow/data/filesystem_dataset_source.py +81 -0
  47. mlflow/data/http_dataset_source.py +145 -0
  48. mlflow/data/huggingface_dataset.py +258 -0
  49. mlflow/data/huggingface_dataset_source.py +118 -0
  50. mlflow/data/meta_dataset.py +104 -0
  51. mlflow/data/numpy_dataset.py +223 -0
  52. mlflow/data/pandas_dataset.py +231 -0
  53. mlflow/data/polars_dataset.py +352 -0
  54. mlflow/data/pyfunc_dataset_mixin.py +31 -0
  55. mlflow/data/schema.py +76 -0
  56. mlflow/data/sources.py +1 -0
  57. mlflow/data/spark_dataset.py +406 -0
  58. mlflow/data/spark_dataset_source.py +74 -0
  59. mlflow/data/spark_delta_utils.py +118 -0
  60. mlflow/data/tensorflow_dataset.py +350 -0
  61. mlflow/data/uc_volume_dataset_source.py +81 -0
  62. mlflow/db.py +27 -0
  63. mlflow/dspy/__init__.py +17 -0
  64. mlflow/dspy/autolog.py +197 -0
  65. mlflow/dspy/callback.py +398 -0
  66. mlflow/dspy/constant.py +1 -0
  67. mlflow/dspy/load.py +93 -0
  68. mlflow/dspy/save.py +393 -0
  69. mlflow/dspy/util.py +109 -0
  70. mlflow/dspy/wrapper.py +226 -0
  71. mlflow/entities/__init__.py +104 -0
  72. mlflow/entities/_mlflow_object.py +52 -0
  73. mlflow/entities/assessment.py +545 -0
  74. mlflow/entities/assessment_error.py +80 -0
  75. mlflow/entities/assessment_source.py +141 -0
  76. mlflow/entities/dataset.py +92 -0
  77. mlflow/entities/dataset_input.py +51 -0
  78. mlflow/entities/dataset_summary.py +62 -0
  79. mlflow/entities/document.py +48 -0
  80. mlflow/entities/experiment.py +109 -0
  81. mlflow/entities/experiment_tag.py +35 -0
  82. mlflow/entities/file_info.py +45 -0
  83. mlflow/entities/input_tag.py +35 -0
  84. mlflow/entities/lifecycle_stage.py +35 -0
  85. mlflow/entities/logged_model.py +228 -0
  86. mlflow/entities/logged_model_input.py +26 -0
  87. mlflow/entities/logged_model_output.py +32 -0
  88. mlflow/entities/logged_model_parameter.py +46 -0
  89. mlflow/entities/logged_model_status.py +74 -0
  90. mlflow/entities/logged_model_tag.py +33 -0
  91. mlflow/entities/metric.py +200 -0
  92. mlflow/entities/model_registry/__init__.py +29 -0
  93. mlflow/entities/model_registry/_model_registry_entity.py +13 -0
  94. mlflow/entities/model_registry/model_version.py +243 -0
  95. mlflow/entities/model_registry/model_version_deployment_job_run_state.py +44 -0
  96. mlflow/entities/model_registry/model_version_deployment_job_state.py +70 -0
  97. mlflow/entities/model_registry/model_version_search.py +25 -0
  98. mlflow/entities/model_registry/model_version_stages.py +25 -0
  99. mlflow/entities/model_registry/model_version_status.py +35 -0
  100. mlflow/entities/model_registry/model_version_tag.py +35 -0
  101. mlflow/entities/model_registry/prompt.py +73 -0
  102. mlflow/entities/model_registry/prompt_version.py +244 -0
  103. mlflow/entities/model_registry/registered_model.py +175 -0
  104. mlflow/entities/model_registry/registered_model_alias.py +35 -0
  105. mlflow/entities/model_registry/registered_model_deployment_job_state.py +39 -0
  106. mlflow/entities/model_registry/registered_model_search.py +25 -0
  107. mlflow/entities/model_registry/registered_model_tag.py +35 -0
  108. mlflow/entities/multipart_upload.py +74 -0
  109. mlflow/entities/param.py +49 -0
  110. mlflow/entities/run.py +97 -0
  111. mlflow/entities/run_data.py +84 -0
  112. mlflow/entities/run_info.py +188 -0
  113. mlflow/entities/run_inputs.py +59 -0
  114. mlflow/entities/run_outputs.py +43 -0
  115. mlflow/entities/run_status.py +41 -0
  116. mlflow/entities/run_tag.py +36 -0
  117. mlflow/entities/source_type.py +31 -0
  118. mlflow/entities/span.py +774 -0
  119. mlflow/entities/span_event.py +96 -0
  120. mlflow/entities/span_status.py +102 -0
  121. mlflow/entities/trace.py +317 -0
  122. mlflow/entities/trace_data.py +71 -0
  123. mlflow/entities/trace_info.py +220 -0
  124. mlflow/entities/trace_info_v2.py +162 -0
  125. mlflow/entities/trace_location.py +173 -0
  126. mlflow/entities/trace_state.py +39 -0
  127. mlflow/entities/trace_status.py +68 -0
  128. mlflow/entities/view_type.py +51 -0
  129. mlflow/environment_variables.py +866 -0
  130. mlflow/evaluation/__init__.py +16 -0
  131. mlflow/evaluation/assessment.py +369 -0
  132. mlflow/evaluation/evaluation.py +411 -0
  133. mlflow/evaluation/evaluation_tag.py +61 -0
  134. mlflow/evaluation/fluent.py +48 -0
  135. mlflow/evaluation/utils.py +201 -0
  136. mlflow/exceptions.py +213 -0
  137. mlflow/experiments.py +140 -0
  138. mlflow/gemini/__init__.py +81 -0
  139. mlflow/gemini/autolog.py +186 -0
  140. mlflow/gemini/chat.py +261 -0
  141. mlflow/genai/__init__.py +71 -0
  142. mlflow/genai/datasets/__init__.py +67 -0
  143. mlflow/genai/datasets/evaluation_dataset.py +131 -0
  144. mlflow/genai/evaluation/__init__.py +3 -0
  145. mlflow/genai/evaluation/base.py +411 -0
  146. mlflow/genai/evaluation/constant.py +23 -0
  147. mlflow/genai/evaluation/utils.py +244 -0
  148. mlflow/genai/judges/__init__.py +21 -0
  149. mlflow/genai/judges/databricks.py +404 -0
  150. mlflow/genai/label_schemas/__init__.py +153 -0
  151. mlflow/genai/label_schemas/label_schemas.py +209 -0
  152. mlflow/genai/labeling/__init__.py +159 -0
  153. mlflow/genai/labeling/labeling.py +250 -0
  154. mlflow/genai/optimize/__init__.py +13 -0
  155. mlflow/genai/optimize/base.py +198 -0
  156. mlflow/genai/optimize/optimizers/__init__.py +4 -0
  157. mlflow/genai/optimize/optimizers/base_optimizer.py +38 -0
  158. mlflow/genai/optimize/optimizers/dspy_mipro_optimizer.py +221 -0
  159. mlflow/genai/optimize/optimizers/dspy_optimizer.py +91 -0
  160. mlflow/genai/optimize/optimizers/utils/dspy_mipro_callback.py +76 -0
  161. mlflow/genai/optimize/optimizers/utils/dspy_mipro_utils.py +18 -0
  162. mlflow/genai/optimize/types.py +75 -0
  163. mlflow/genai/optimize/util.py +30 -0
  164. mlflow/genai/prompts/__init__.py +206 -0
  165. mlflow/genai/scheduled_scorers.py +431 -0
  166. mlflow/genai/scorers/__init__.py +26 -0
  167. mlflow/genai/scorers/base.py +492 -0
  168. mlflow/genai/scorers/builtin_scorers.py +765 -0
  169. mlflow/genai/scorers/scorer_utils.py +138 -0
  170. mlflow/genai/scorers/validation.py +165 -0
  171. mlflow/genai/utils/data_validation.py +146 -0
  172. mlflow/genai/utils/enum_utils.py +23 -0
  173. mlflow/genai/utils/trace_utils.py +211 -0
  174. mlflow/groq/__init__.py +42 -0
  175. mlflow/groq/_groq_autolog.py +74 -0
  176. mlflow/johnsnowlabs/__init__.py +888 -0
  177. mlflow/langchain/__init__.py +24 -0
  178. mlflow/langchain/api_request_parallel_processor.py +330 -0
  179. mlflow/langchain/autolog.py +147 -0
  180. mlflow/langchain/chat_agent_langgraph.py +340 -0
  181. mlflow/langchain/constant.py +1 -0
  182. mlflow/langchain/constants.py +1 -0
  183. mlflow/langchain/databricks_dependencies.py +444 -0
  184. mlflow/langchain/langchain_tracer.py +597 -0
  185. mlflow/langchain/model.py +919 -0
  186. mlflow/langchain/output_parsers.py +142 -0
  187. mlflow/langchain/retriever_chain.py +153 -0
  188. mlflow/langchain/runnables.py +527 -0
  189. mlflow/langchain/utils/chat.py +402 -0
  190. mlflow/langchain/utils/logging.py +671 -0
  191. mlflow/langchain/utils/serialization.py +36 -0
  192. mlflow/legacy_databricks_cli/__init__.py +0 -0
  193. mlflow/legacy_databricks_cli/configure/__init__.py +0 -0
  194. mlflow/legacy_databricks_cli/configure/provider.py +482 -0
  195. mlflow/litellm/__init__.py +175 -0
  196. mlflow/llama_index/__init__.py +22 -0
  197. mlflow/llama_index/autolog.py +55 -0
  198. mlflow/llama_index/chat.py +43 -0
  199. mlflow/llama_index/constant.py +1 -0
  200. mlflow/llama_index/model.py +577 -0
  201. mlflow/llama_index/pyfunc_wrapper.py +332 -0
  202. mlflow/llama_index/serialize_objects.py +188 -0
  203. mlflow/llama_index/tracer.py +561 -0
  204. mlflow/metrics/__init__.py +479 -0
  205. mlflow/metrics/base.py +39 -0
  206. mlflow/metrics/genai/__init__.py +25 -0
  207. mlflow/metrics/genai/base.py +101 -0
  208. mlflow/metrics/genai/genai_metric.py +771 -0
  209. mlflow/metrics/genai/metric_definitions.py +450 -0
  210. mlflow/metrics/genai/model_utils.py +371 -0
  211. mlflow/metrics/genai/prompt_template.py +68 -0
  212. mlflow/metrics/genai/prompts/__init__.py +0 -0
  213. mlflow/metrics/genai/prompts/v1.py +422 -0
  214. mlflow/metrics/genai/utils.py +6 -0
  215. mlflow/metrics/metric_definitions.py +619 -0
  216. mlflow/mismatch.py +34 -0
  217. mlflow/mistral/__init__.py +34 -0
  218. mlflow/mistral/autolog.py +71 -0
  219. mlflow/mistral/chat.py +135 -0
  220. mlflow/ml_package_versions.py +452 -0
  221. mlflow/models/__init__.py +97 -0
  222. mlflow/models/auth_policy.py +83 -0
  223. mlflow/models/cli.py +354 -0
  224. mlflow/models/container/__init__.py +294 -0
  225. mlflow/models/container/scoring_server/__init__.py +0 -0
  226. mlflow/models/container/scoring_server/nginx.conf +39 -0
  227. mlflow/models/dependencies_schemas.py +287 -0
  228. mlflow/models/display_utils.py +158 -0
  229. mlflow/models/docker_utils.py +211 -0
  230. mlflow/models/evaluation/__init__.py +23 -0
  231. mlflow/models/evaluation/_shap_patch.py +64 -0
  232. mlflow/models/evaluation/artifacts.py +194 -0
  233. mlflow/models/evaluation/base.py +1811 -0
  234. mlflow/models/evaluation/calibration_curve.py +109 -0
  235. mlflow/models/evaluation/default_evaluator.py +996 -0
  236. mlflow/models/evaluation/deprecated.py +23 -0
  237. mlflow/models/evaluation/evaluator_registry.py +80 -0
  238. mlflow/models/evaluation/evaluators/classifier.py +704 -0
  239. mlflow/models/evaluation/evaluators/default.py +233 -0
  240. mlflow/models/evaluation/evaluators/regressor.py +96 -0
  241. mlflow/models/evaluation/evaluators/shap.py +296 -0
  242. mlflow/models/evaluation/lift_curve.py +178 -0
  243. mlflow/models/evaluation/utils/metric.py +123 -0
  244. mlflow/models/evaluation/utils/trace.py +179 -0
  245. mlflow/models/evaluation/validation.py +434 -0
  246. mlflow/models/flavor_backend.py +93 -0
  247. mlflow/models/flavor_backend_registry.py +53 -0
  248. mlflow/models/model.py +1639 -0
  249. mlflow/models/model_config.py +150 -0
  250. mlflow/models/notebook_resources/agent_evaluation_template.html +235 -0
  251. mlflow/models/notebook_resources/eval_with_dataset_example.py +22 -0
  252. mlflow/models/notebook_resources/eval_with_synthetic_example.py +22 -0
  253. mlflow/models/python_api.py +369 -0
  254. mlflow/models/rag_signatures.py +128 -0
  255. mlflow/models/resources.py +321 -0
  256. mlflow/models/signature.py +662 -0
  257. mlflow/models/utils.py +2054 -0
  258. mlflow/models/wheeled_model.py +280 -0
  259. mlflow/openai/__init__.py +57 -0
  260. mlflow/openai/_agent_tracer.py +364 -0
  261. mlflow/openai/api_request_parallel_processor.py +131 -0
  262. mlflow/openai/autolog.py +509 -0
  263. mlflow/openai/constant.py +1 -0
  264. mlflow/openai/model.py +824 -0
  265. mlflow/openai/utils/chat_schema.py +367 -0
  266. mlflow/optuna/__init__.py +3 -0
  267. mlflow/optuna/storage.py +646 -0
  268. mlflow/plugins/__init__.py +72 -0
  269. mlflow/plugins/base.py +358 -0
  270. mlflow/plugins/builtin/__init__.py +24 -0
  271. mlflow/plugins/builtin/pytorch_plugin.py +150 -0
  272. mlflow/plugins/builtin/sklearn_plugin.py +158 -0
  273. mlflow/plugins/builtin/transformers_plugin.py +187 -0
  274. mlflow/plugins/cli.py +321 -0
  275. mlflow/plugins/discovery.py +340 -0
  276. mlflow/plugins/manager.py +465 -0
  277. mlflow/plugins/registry.py +316 -0
  278. mlflow/plugins/templates/framework_plugin_template.py +329 -0
  279. mlflow/prompt/constants.py +20 -0
  280. mlflow/prompt/promptlab_model.py +197 -0
  281. mlflow/prompt/registry_utils.py +248 -0
  282. mlflow/promptflow/__init__.py +495 -0
  283. mlflow/protos/__init__.py +0 -0
  284. mlflow/protos/assessments_pb2.py +174 -0
  285. mlflow/protos/databricks_artifacts_pb2.py +489 -0
  286. mlflow/protos/databricks_filesystem_service_pb2.py +196 -0
  287. mlflow/protos/databricks_managed_catalog_messages_pb2.py +95 -0
  288. mlflow/protos/databricks_managed_catalog_service_pb2.py +86 -0
  289. mlflow/protos/databricks_pb2.py +267 -0
  290. mlflow/protos/databricks_trace_server_pb2.py +374 -0
  291. mlflow/protos/databricks_uc_registry_messages_pb2.py +1249 -0
  292. mlflow/protos/databricks_uc_registry_service_pb2.py +170 -0
  293. mlflow/protos/facet_feature_statistics_pb2.py +296 -0
  294. mlflow/protos/internal_pb2.py +77 -0
  295. mlflow/protos/mlflow_artifacts_pb2.py +336 -0
  296. mlflow/protos/model_registry_pb2.py +1073 -0
  297. mlflow/protos/scalapb/__init__.py +0 -0
  298. mlflow/protos/scalapb/scalapb_pb2.py +104 -0
  299. mlflow/protos/service_pb2.py +2600 -0
  300. mlflow/protos/unity_catalog_oss_messages_pb2.py +457 -0
  301. mlflow/protos/unity_catalog_oss_service_pb2.py +130 -0
  302. mlflow/protos/unity_catalog_prompt_messages_pb2.py +447 -0
  303. mlflow/protos/unity_catalog_prompt_messages_pb2_grpc.py +24 -0
  304. mlflow/protos/unity_catalog_prompt_service_pb2.py +164 -0
  305. mlflow/protos/unity_catalog_prompt_service_pb2_grpc.py +785 -0
  306. mlflow/py.typed +0 -0
  307. mlflow/pydantic_ai/__init__.py +57 -0
  308. mlflow/pydantic_ai/autolog.py +173 -0
  309. mlflow/pyfunc/__init__.py +3844 -0
  310. mlflow/pyfunc/_mlflow_pyfunc_backend_predict.py +61 -0
  311. mlflow/pyfunc/backend.py +523 -0
  312. mlflow/pyfunc/context.py +78 -0
  313. mlflow/pyfunc/dbconnect_artifact_cache.py +144 -0
  314. mlflow/pyfunc/loaders/__init__.py +7 -0
  315. mlflow/pyfunc/loaders/chat_agent.py +117 -0
  316. mlflow/pyfunc/loaders/chat_model.py +125 -0
  317. mlflow/pyfunc/loaders/code_model.py +31 -0
  318. mlflow/pyfunc/loaders/responses_agent.py +112 -0
  319. mlflow/pyfunc/mlserver.py +46 -0
  320. mlflow/pyfunc/model.py +1473 -0
  321. mlflow/pyfunc/scoring_server/__init__.py +604 -0
  322. mlflow/pyfunc/scoring_server/app.py +7 -0
  323. mlflow/pyfunc/scoring_server/client.py +146 -0
  324. mlflow/pyfunc/spark_model_cache.py +48 -0
  325. mlflow/pyfunc/stdin_server.py +44 -0
  326. mlflow/pyfunc/utils/__init__.py +3 -0
  327. mlflow/pyfunc/utils/data_validation.py +224 -0
  328. mlflow/pyfunc/utils/environment.py +22 -0
  329. mlflow/pyfunc/utils/input_converter.py +47 -0
  330. mlflow/pyfunc/utils/serving_data_parser.py +11 -0
  331. mlflow/pytorch/__init__.py +1171 -0
  332. mlflow/pytorch/_lightning_autolog.py +580 -0
  333. mlflow/pytorch/_pytorch_autolog.py +50 -0
  334. mlflow/pytorch/pickle_module.py +35 -0
  335. mlflow/rfunc/__init__.py +42 -0
  336. mlflow/rfunc/backend.py +134 -0
  337. mlflow/runs.py +89 -0
  338. mlflow/server/__init__.py +302 -0
  339. mlflow/server/auth/__init__.py +1224 -0
  340. mlflow/server/auth/__main__.py +4 -0
  341. mlflow/server/auth/basic_auth.ini +6 -0
  342. mlflow/server/auth/cli.py +11 -0
  343. mlflow/server/auth/client.py +537 -0
  344. mlflow/server/auth/config.py +34 -0
  345. mlflow/server/auth/db/__init__.py +0 -0
  346. mlflow/server/auth/db/cli.py +18 -0
  347. mlflow/server/auth/db/migrations/__init__.py +0 -0
  348. mlflow/server/auth/db/migrations/alembic.ini +110 -0
  349. mlflow/server/auth/db/migrations/env.py +76 -0
  350. mlflow/server/auth/db/migrations/versions/8606fa83a998_initial_migration.py +51 -0
  351. mlflow/server/auth/db/migrations/versions/__init__.py +0 -0
  352. mlflow/server/auth/db/models.py +67 -0
  353. mlflow/server/auth/db/utils.py +37 -0
  354. mlflow/server/auth/entities.py +165 -0
  355. mlflow/server/auth/logo.py +14 -0
  356. mlflow/server/auth/permissions.py +65 -0
  357. mlflow/server/auth/routes.py +18 -0
  358. mlflow/server/auth/sqlalchemy_store.py +263 -0
  359. mlflow/server/graphql/__init__.py +0 -0
  360. mlflow/server/graphql/autogenerated_graphql_schema.py +353 -0
  361. mlflow/server/graphql/graphql_custom_scalars.py +24 -0
  362. mlflow/server/graphql/graphql_errors.py +15 -0
  363. mlflow/server/graphql/graphql_no_batching.py +89 -0
  364. mlflow/server/graphql/graphql_schema_extensions.py +74 -0
  365. mlflow/server/handlers.py +3217 -0
  366. mlflow/server/prometheus_exporter.py +17 -0
  367. mlflow/server/validation.py +30 -0
  368. mlflow/shap/__init__.py +691 -0
  369. mlflow/sklearn/__init__.py +1994 -0
  370. mlflow/sklearn/utils.py +1041 -0
  371. mlflow/smolagents/__init__.py +66 -0
  372. mlflow/smolagents/autolog.py +139 -0
  373. mlflow/smolagents/chat.py +29 -0
  374. mlflow/store/__init__.py +10 -0
  375. mlflow/store/_unity_catalog/__init__.py +1 -0
  376. mlflow/store/_unity_catalog/lineage/__init__.py +1 -0
  377. mlflow/store/_unity_catalog/lineage/constants.py +2 -0
  378. mlflow/store/_unity_catalog/registry/__init__.py +6 -0
  379. mlflow/store/_unity_catalog/registry/prompt_info.py +75 -0
  380. mlflow/store/_unity_catalog/registry/rest_store.py +1740 -0
  381. mlflow/store/_unity_catalog/registry/uc_oss_rest_store.py +507 -0
  382. mlflow/store/_unity_catalog/registry/utils.py +121 -0
  383. mlflow/store/artifact/__init__.py +0 -0
  384. mlflow/store/artifact/artifact_repo.py +472 -0
  385. mlflow/store/artifact/artifact_repository_registry.py +154 -0
  386. mlflow/store/artifact/azure_blob_artifact_repo.py +275 -0
  387. mlflow/store/artifact/azure_data_lake_artifact_repo.py +295 -0
  388. mlflow/store/artifact/cli.py +141 -0
  389. mlflow/store/artifact/cloud_artifact_repo.py +332 -0
  390. mlflow/store/artifact/databricks_artifact_repo.py +729 -0
  391. mlflow/store/artifact/databricks_artifact_repo_resources.py +301 -0
  392. mlflow/store/artifact/databricks_logged_model_artifact_repo.py +93 -0
  393. mlflow/store/artifact/databricks_models_artifact_repo.py +216 -0
  394. mlflow/store/artifact/databricks_sdk_artifact_repo.py +134 -0
  395. mlflow/store/artifact/databricks_sdk_models_artifact_repo.py +97 -0
  396. mlflow/store/artifact/dbfs_artifact_repo.py +240 -0
  397. mlflow/store/artifact/ftp_artifact_repo.py +132 -0
  398. mlflow/store/artifact/gcs_artifact_repo.py +296 -0
  399. mlflow/store/artifact/hdfs_artifact_repo.py +209 -0
  400. mlflow/store/artifact/http_artifact_repo.py +218 -0
  401. mlflow/store/artifact/local_artifact_repo.py +142 -0
  402. mlflow/store/artifact/mlflow_artifacts_repo.py +94 -0
  403. mlflow/store/artifact/models_artifact_repo.py +259 -0
  404. mlflow/store/artifact/optimized_s3_artifact_repo.py +356 -0
  405. mlflow/store/artifact/presigned_url_artifact_repo.py +173 -0
  406. mlflow/store/artifact/r2_artifact_repo.py +70 -0
  407. mlflow/store/artifact/runs_artifact_repo.py +265 -0
  408. mlflow/store/artifact/s3_artifact_repo.py +330 -0
  409. mlflow/store/artifact/sftp_artifact_repo.py +141 -0
  410. mlflow/store/artifact/uc_volume_artifact_repo.py +76 -0
  411. mlflow/store/artifact/unity_catalog_models_artifact_repo.py +168 -0
  412. mlflow/store/artifact/unity_catalog_oss_models_artifact_repo.py +168 -0
  413. mlflow/store/artifact/utils/__init__.py +0 -0
  414. mlflow/store/artifact/utils/models.py +148 -0
  415. mlflow/store/db/__init__.py +0 -0
  416. mlflow/store/db/base_sql_model.py +3 -0
  417. mlflow/store/db/db_types.py +10 -0
  418. mlflow/store/db/utils.py +314 -0
  419. mlflow/store/db_migrations/__init__.py +0 -0
  420. mlflow/store/db_migrations/alembic.ini +74 -0
  421. mlflow/store/db_migrations/env.py +84 -0
  422. mlflow/store/db_migrations/versions/0584bdc529eb_add_cascading_deletion_to_datasets_from_experiments.py +88 -0
  423. mlflow/store/db_migrations/versions/0a8213491aaa_drop_duplicate_killed_constraint.py +49 -0
  424. mlflow/store/db_migrations/versions/0c779009ac13_add_deleted_time_field_to_runs_table.py +24 -0
  425. mlflow/store/db_migrations/versions/181f10493468_allow_nulls_for_metric_values.py +35 -0
  426. mlflow/store/db_migrations/versions/27a6a02d2cf1_add_model_version_tags_table.py +38 -0
  427. mlflow/store/db_migrations/versions/2b4d017a5e9b_add_model_registry_tables_to_db.py +77 -0
  428. mlflow/store/db_migrations/versions/2d6e25af4d3e_increase_max_param_val_length.py +33 -0
  429. mlflow/store/db_migrations/versions/3500859a5d39_add_model_aliases_table.py +50 -0
  430. mlflow/store/db_migrations/versions/39d1c3be5f05_add_is_nan_constraint_for_metrics_tables_if_necessary.py +41 -0
  431. mlflow/store/db_migrations/versions/400f98739977_add_logged_model_tables.py +123 -0
  432. mlflow/store/db_migrations/versions/4465047574b1_increase_max_dataset_schema_size.py +38 -0
  433. mlflow/store/db_migrations/versions/451aebb31d03_add_metric_step.py +35 -0
  434. mlflow/store/db_migrations/versions/5b0e9adcef9c_add_cascade_deletion_to_trace_tables_fk.py +40 -0
  435. mlflow/store/db_migrations/versions/6953534de441_add_step_to_inputs_table.py +25 -0
  436. mlflow/store/db_migrations/versions/728d730b5ebd_add_registered_model_tags_table.py +38 -0
  437. mlflow/store/db_migrations/versions/7ac759974ad8_update_run_tags_with_larger_limit.py +36 -0
  438. mlflow/store/db_migrations/versions/7f2a7d5fae7d_add_datasets_inputs_input_tags_tables.py +82 -0
  439. mlflow/store/db_migrations/versions/84291f40a231_add_run_link_to_model_version.py +26 -0
  440. mlflow/store/db_migrations/versions/867495a8f9d4_add_trace_tables.py +90 -0
  441. mlflow/store/db_migrations/versions/89d4b8295536_create_latest_metrics_table.py +169 -0
  442. mlflow/store/db_migrations/versions/90e64c465722_migrate_user_column_to_tags.py +64 -0
  443. mlflow/store/db_migrations/versions/97727af70f4d_creation_time_last_update_time_experiments.py +25 -0
  444. mlflow/store/db_migrations/versions/__init__.py +0 -0
  445. mlflow/store/db_migrations/versions/a8c4a736bde6_allow_nulls_for_run_id.py +27 -0
  446. mlflow/store/db_migrations/versions/acf3f17fdcc7_add_storage_location_field_to_model_.py +29 -0
  447. mlflow/store/db_migrations/versions/bd07f7e963c5_create_index_on_run_uuid.py +26 -0
  448. mlflow/store/db_migrations/versions/bda7b8c39065_increase_model_version_tag_value_limit.py +38 -0
  449. mlflow/store/db_migrations/versions/c48cb773bb87_reset_default_value_for_is_nan_in_metrics_table_for_mysql.py +41 -0
  450. mlflow/store/db_migrations/versions/cbc13b556ace_add_v3_trace_schema_columns.py +31 -0
  451. mlflow/store/db_migrations/versions/cc1f77228345_change_param_value_length_to_500.py +34 -0
  452. mlflow/store/db_migrations/versions/cfd24bdc0731_update_run_status_constraint_with_killed.py +78 -0
  453. mlflow/store/db_migrations/versions/df50e92ffc5e_add_experiment_tags_table.py +38 -0
  454. mlflow/store/db_migrations/versions/f5a4f2784254_increase_run_tag_value_limit.py +36 -0
  455. mlflow/store/entities/__init__.py +3 -0
  456. mlflow/store/entities/paged_list.py +18 -0
  457. mlflow/store/model_registry/__init__.py +10 -0
  458. mlflow/store/model_registry/abstract_store.py +1081 -0
  459. mlflow/store/model_registry/base_rest_store.py +44 -0
  460. mlflow/store/model_registry/databricks_workspace_model_registry_rest_store.py +37 -0
  461. mlflow/store/model_registry/dbmodels/__init__.py +0 -0
  462. mlflow/store/model_registry/dbmodels/models.py +206 -0
  463. mlflow/store/model_registry/file_store.py +1091 -0
  464. mlflow/store/model_registry/rest_store.py +481 -0
  465. mlflow/store/model_registry/sqlalchemy_store.py +1286 -0
  466. mlflow/store/tracking/__init__.py +23 -0
  467. mlflow/store/tracking/abstract_store.py +816 -0
  468. mlflow/store/tracking/dbmodels/__init__.py +0 -0
  469. mlflow/store/tracking/dbmodels/initial_models.py +243 -0
  470. mlflow/store/tracking/dbmodels/models.py +1073 -0
  471. mlflow/store/tracking/file_store.py +2438 -0
  472. mlflow/store/tracking/postgres_managed_identity.py +146 -0
  473. mlflow/store/tracking/rest_store.py +1131 -0
  474. mlflow/store/tracking/sqlalchemy_store.py +2785 -0
  475. mlflow/system_metrics/__init__.py +61 -0
  476. mlflow/system_metrics/metrics/__init__.py +0 -0
  477. mlflow/system_metrics/metrics/base_metrics_monitor.py +32 -0
  478. mlflow/system_metrics/metrics/cpu_monitor.py +23 -0
  479. mlflow/system_metrics/metrics/disk_monitor.py +21 -0
  480. mlflow/system_metrics/metrics/gpu_monitor.py +71 -0
  481. mlflow/system_metrics/metrics/network_monitor.py +34 -0
  482. mlflow/system_metrics/metrics/rocm_monitor.py +123 -0
  483. mlflow/system_metrics/system_metrics_monitor.py +198 -0
  484. mlflow/tracing/__init__.py +16 -0
  485. mlflow/tracing/assessment.py +356 -0
  486. mlflow/tracing/client.py +531 -0
  487. mlflow/tracing/config.py +125 -0
  488. mlflow/tracing/constant.py +105 -0
  489. mlflow/tracing/destination.py +81 -0
  490. mlflow/tracing/display/__init__.py +40 -0
  491. mlflow/tracing/display/display_handler.py +196 -0
  492. mlflow/tracing/export/async_export_queue.py +186 -0
  493. mlflow/tracing/export/inference_table.py +138 -0
  494. mlflow/tracing/export/mlflow_v3.py +137 -0
  495. mlflow/tracing/export/utils.py +70 -0
  496. mlflow/tracing/fluent.py +1417 -0
  497. mlflow/tracing/processor/base_mlflow.py +199 -0
  498. mlflow/tracing/processor/inference_table.py +175 -0
  499. mlflow/tracing/processor/mlflow_v3.py +47 -0
  500. mlflow/tracing/processor/otel.py +73 -0
  501. mlflow/tracing/provider.py +487 -0
  502. mlflow/tracing/trace_manager.py +200 -0
  503. mlflow/tracing/utils/__init__.py +616 -0
  504. mlflow/tracing/utils/artifact_utils.py +28 -0
  505. mlflow/tracing/utils/copy.py +55 -0
  506. mlflow/tracing/utils/environment.py +55 -0
  507. mlflow/tracing/utils/exception.py +21 -0
  508. mlflow/tracing/utils/once.py +35 -0
  509. mlflow/tracing/utils/otlp.py +63 -0
  510. mlflow/tracing/utils/processor.py +54 -0
  511. mlflow/tracing/utils/search.py +292 -0
  512. mlflow/tracing/utils/timeout.py +250 -0
  513. mlflow/tracing/utils/token.py +19 -0
  514. mlflow/tracing/utils/truncation.py +124 -0
  515. mlflow/tracing/utils/warning.py +76 -0
  516. mlflow/tracking/__init__.py +39 -0
  517. mlflow/tracking/_model_registry/__init__.py +1 -0
  518. mlflow/tracking/_model_registry/client.py +764 -0
  519. mlflow/tracking/_model_registry/fluent.py +853 -0
  520. mlflow/tracking/_model_registry/registry.py +67 -0
  521. mlflow/tracking/_model_registry/utils.py +251 -0
  522. mlflow/tracking/_tracking_service/__init__.py +0 -0
  523. mlflow/tracking/_tracking_service/client.py +883 -0
  524. mlflow/tracking/_tracking_service/registry.py +56 -0
  525. mlflow/tracking/_tracking_service/utils.py +275 -0
  526. mlflow/tracking/artifact_utils.py +179 -0
  527. mlflow/tracking/client.py +5900 -0
  528. mlflow/tracking/context/__init__.py +0 -0
  529. mlflow/tracking/context/abstract_context.py +35 -0
  530. mlflow/tracking/context/databricks_cluster_context.py +15 -0
  531. mlflow/tracking/context/databricks_command_context.py +15 -0
  532. mlflow/tracking/context/databricks_job_context.py +49 -0
  533. mlflow/tracking/context/databricks_notebook_context.py +41 -0
  534. mlflow/tracking/context/databricks_repo_context.py +43 -0
  535. mlflow/tracking/context/default_context.py +51 -0
  536. mlflow/tracking/context/git_context.py +32 -0
  537. mlflow/tracking/context/registry.py +98 -0
  538. mlflow/tracking/context/system_environment_context.py +15 -0
  539. mlflow/tracking/default_experiment/__init__.py +1 -0
  540. mlflow/tracking/default_experiment/abstract_context.py +43 -0
  541. mlflow/tracking/default_experiment/databricks_notebook_experiment_provider.py +44 -0
  542. mlflow/tracking/default_experiment/registry.py +75 -0
  543. mlflow/tracking/fluent.py +3595 -0
  544. mlflow/tracking/metric_value_conversion_utils.py +93 -0
  545. mlflow/tracking/multimedia.py +206 -0
  546. mlflow/tracking/registry.py +86 -0
  547. mlflow/tracking/request_auth/__init__.py +0 -0
  548. mlflow/tracking/request_auth/abstract_request_auth_provider.py +34 -0
  549. mlflow/tracking/request_auth/registry.py +60 -0
  550. mlflow/tracking/request_header/__init__.py +0 -0
  551. mlflow/tracking/request_header/abstract_request_header_provider.py +36 -0
  552. mlflow/tracking/request_header/databricks_request_header_provider.py +38 -0
  553. mlflow/tracking/request_header/default_request_header_provider.py +17 -0
  554. mlflow/tracking/request_header/registry.py +79 -0
  555. mlflow/transformers/__init__.py +2982 -0
  556. mlflow/transformers/flavor_config.py +258 -0
  557. mlflow/transformers/hub_utils.py +83 -0
  558. mlflow/transformers/llm_inference_utils.py +468 -0
  559. mlflow/transformers/model_io.py +301 -0
  560. mlflow/transformers/peft.py +51 -0
  561. mlflow/transformers/signature.py +183 -0
  562. mlflow/transformers/torch_utils.py +55 -0
  563. mlflow/types/__init__.py +21 -0
  564. mlflow/types/agent.py +270 -0
  565. mlflow/types/chat.py +240 -0
  566. mlflow/types/llm.py +935 -0
  567. mlflow/types/responses.py +139 -0
  568. mlflow/types/responses_helpers.py +416 -0
  569. mlflow/types/schema.py +1505 -0
  570. mlflow/types/type_hints.py +647 -0
  571. mlflow/types/utils.py +753 -0
  572. mlflow/utils/__init__.py +283 -0
  573. mlflow/utils/_capture_modules.py +256 -0
  574. mlflow/utils/_capture_transformers_modules.py +75 -0
  575. mlflow/utils/_spark_utils.py +201 -0
  576. mlflow/utils/_unity_catalog_oss_utils.py +97 -0
  577. mlflow/utils/_unity_catalog_utils.py +479 -0
  578. mlflow/utils/annotations.py +218 -0
  579. mlflow/utils/arguments_utils.py +16 -0
  580. mlflow/utils/async_logging/__init__.py +1 -0
  581. mlflow/utils/async_logging/async_artifacts_logging_queue.py +258 -0
  582. mlflow/utils/async_logging/async_logging_queue.py +366 -0
  583. mlflow/utils/async_logging/run_artifact.py +38 -0
  584. mlflow/utils/async_logging/run_batch.py +58 -0
  585. mlflow/utils/async_logging/run_operations.py +49 -0
  586. mlflow/utils/autologging_utils/__init__.py +737 -0
  587. mlflow/utils/autologging_utils/client.py +432 -0
  588. mlflow/utils/autologging_utils/config.py +33 -0
  589. mlflow/utils/autologging_utils/events.py +294 -0
  590. mlflow/utils/autologging_utils/logging_and_warnings.py +328 -0
  591. mlflow/utils/autologging_utils/metrics_queue.py +71 -0
  592. mlflow/utils/autologging_utils/safety.py +1104 -0
  593. mlflow/utils/autologging_utils/versioning.py +95 -0
  594. mlflow/utils/checkpoint_utils.py +206 -0
  595. mlflow/utils/class_utils.py +6 -0
  596. mlflow/utils/cli_args.py +257 -0
  597. mlflow/utils/conda.py +354 -0
  598. mlflow/utils/credentials.py +231 -0
  599. mlflow/utils/data_utils.py +17 -0
  600. mlflow/utils/databricks_utils.py +1436 -0
  601. mlflow/utils/docstring_utils.py +477 -0
  602. mlflow/utils/doctor.py +133 -0
  603. mlflow/utils/download_cloud_file_chunk.py +43 -0
  604. mlflow/utils/env_manager.py +16 -0
  605. mlflow/utils/env_pack.py +131 -0
  606. mlflow/utils/environment.py +1009 -0
  607. mlflow/utils/exception_utils.py +14 -0
  608. mlflow/utils/file_utils.py +978 -0
  609. mlflow/utils/git_utils.py +77 -0
  610. mlflow/utils/gorilla.py +797 -0
  611. mlflow/utils/import_hooks/__init__.py +363 -0
  612. mlflow/utils/lazy_load.py +51 -0
  613. mlflow/utils/logging_utils.py +168 -0
  614. mlflow/utils/mime_type_utils.py +58 -0
  615. mlflow/utils/mlflow_tags.py +103 -0
  616. mlflow/utils/model_utils.py +486 -0
  617. mlflow/utils/name_utils.py +346 -0
  618. mlflow/utils/nfs_on_spark.py +62 -0
  619. mlflow/utils/openai_utils.py +164 -0
  620. mlflow/utils/os.py +12 -0
  621. mlflow/utils/oss_registry_utils.py +29 -0
  622. mlflow/utils/plugins.py +17 -0
  623. mlflow/utils/process.py +182 -0
  624. mlflow/utils/promptlab_utils.py +146 -0
  625. mlflow/utils/proto_json_utils.py +743 -0
  626. mlflow/utils/pydantic_utils.py +54 -0
  627. mlflow/utils/request_utils.py +279 -0
  628. mlflow/utils/requirements_utils.py +704 -0
  629. mlflow/utils/rest_utils.py +673 -0
  630. mlflow/utils/search_logged_model_utils.py +127 -0
  631. mlflow/utils/search_utils.py +2111 -0
  632. mlflow/utils/secure_loading.py +221 -0
  633. mlflow/utils/security_validation.py +384 -0
  634. mlflow/utils/server_cli_utils.py +61 -0
  635. mlflow/utils/spark_utils.py +15 -0
  636. mlflow/utils/string_utils.py +138 -0
  637. mlflow/utils/thread_utils.py +63 -0
  638. mlflow/utils/time.py +54 -0
  639. mlflow/utils/timeout.py +42 -0
  640. mlflow/utils/uri.py +572 -0
  641. mlflow/utils/validation.py +662 -0
  642. mlflow/utils/virtualenv.py +458 -0
  643. mlflow/utils/warnings_utils.py +25 -0
  644. mlflow/utils/yaml_utils.py +179 -0
  645. mlflow/version.py +24 -0
@@ -0,0 +1,1073 @@
1
+ import sqlalchemy as sa
2
+ from sqlalchemy import (
3
+ BigInteger,
4
+ Boolean,
5
+ CheckConstraint,
6
+ Column,
7
+ ForeignKey,
8
+ ForeignKeyConstraint,
9
+ Index,
10
+ Integer,
11
+ PrimaryKeyConstraint,
12
+ String,
13
+ Text,
14
+ UnicodeText,
15
+ )
16
+ from sqlalchemy.orm import backref, relationship
17
+
18
+ from mlflow.entities import (
19
+ Dataset,
20
+ Experiment,
21
+ ExperimentTag,
22
+ InputTag,
23
+ Metric,
24
+ Param,
25
+ Run,
26
+ RunData,
27
+ RunInfo,
28
+ RunStatus,
29
+ RunTag,
30
+ SourceType,
31
+ TraceInfo,
32
+ ViewType,
33
+ )
34
+ from mlflow.entities.lifecycle_stage import LifecycleStage
35
+ from mlflow.entities.logged_model import LoggedModel
36
+ from mlflow.entities.logged_model_parameter import LoggedModelParameter
37
+ from mlflow.entities.logged_model_status import LoggedModelStatus
38
+ from mlflow.entities.logged_model_tag import LoggedModelTag
39
+ from mlflow.entities.trace_location import TraceLocation
40
+ from mlflow.entities.trace_state import TraceState
41
+ from mlflow.store.db.base_sql_model import Base
42
+ from mlflow.utils.mlflow_tags import _get_run_name_from_tags
43
+ from mlflow.utils.time import get_current_time_millis
44
+
45
+ SourceTypes = [
46
+ SourceType.to_string(SourceType.NOTEBOOK),
47
+ SourceType.to_string(SourceType.JOB),
48
+ SourceType.to_string(SourceType.LOCAL),
49
+ SourceType.to_string(SourceType.UNKNOWN),
50
+ SourceType.to_string(SourceType.PROJECT),
51
+ ]
52
+
53
+ RunStatusTypes = [
54
+ RunStatus.to_string(RunStatus.SCHEDULED),
55
+ RunStatus.to_string(RunStatus.FAILED),
56
+ RunStatus.to_string(RunStatus.FINISHED),
57
+ RunStatus.to_string(RunStatus.RUNNING),
58
+ RunStatus.to_string(RunStatus.KILLED),
59
+ ]
60
+
61
+
62
+ class SqlExperiment(Base):
63
+ """
64
+ DB model for :py:class:`mlflow.entities.Experiment`. These are recorded in ``experiment`` table.
65
+ """
66
+
67
+ __tablename__ = "experiments"
68
+
69
+ experiment_id = Column(Integer, autoincrement=True)
70
+ """
71
+ Experiment ID: `Integer`. *Primary Key* for ``experiment`` table.
72
+ """
73
+ name = Column(String(256), unique=True, nullable=False)
74
+ """
75
+ Experiment name: `String` (limit 256 characters). Defined as *Unique* and *Non null* in
76
+ table schema.
77
+ """
78
+ artifact_location = Column(String(256), nullable=True)
79
+ """
80
+ Default artifact location for this experiment: `String` (limit 256 characters). Defined as
81
+ *Non null* in table schema.
82
+ """
83
+ lifecycle_stage = Column(String(32), default=LifecycleStage.ACTIVE)
84
+ """
85
+ Lifecycle Stage of experiment: `String` (limit 32 characters).
86
+ Can be either ``active`` (default) or ``deleted``.
87
+ """
88
+ creation_time = Column(BigInteger(), default=get_current_time_millis)
89
+ """
90
+ Creation time of experiment: `BigInteger`.
91
+ """
92
+ last_update_time = Column(BigInteger(), default=get_current_time_millis)
93
+ """
94
+ Last Update time of experiment: `BigInteger`.
95
+ """
96
+
97
+ __table_args__ = (
98
+ CheckConstraint(
99
+ lifecycle_stage.in_(LifecycleStage.view_type_to_stages(ViewType.ALL)),
100
+ name="experiments_lifecycle_stage",
101
+ ),
102
+ PrimaryKeyConstraint("experiment_id", name="experiment_pk"),
103
+ )
104
+
105
+ def __repr__(self):
106
+ return f"<SqlExperiment ({self.experiment_id}, {self.name})>"
107
+
108
+ def to_mlflow_entity(self):
109
+ """
110
+ Convert DB model to corresponding MLflow entity.
111
+
112
+ Returns:
113
+ :py:class:`mlflow.entities.Experiment`.
114
+ """
115
+ return Experiment(
116
+ experiment_id=str(self.experiment_id),
117
+ name=self.name,
118
+ artifact_location=self.artifact_location,
119
+ lifecycle_stage=self.lifecycle_stage,
120
+ tags=[t.to_mlflow_entity() for t in self.tags],
121
+ creation_time=self.creation_time,
122
+ last_update_time=self.last_update_time,
123
+ )
124
+
125
+
126
+ class SqlRun(Base):
127
+ """
128
+ DB model for :py:class:`mlflow.entities.Run`. These are recorded in ``runs`` table.
129
+ """
130
+
131
+ __tablename__ = "runs"
132
+
133
+ run_uuid = Column(String(32), nullable=False)
134
+ """
135
+ Run UUID: `String` (limit 32 characters). *Primary Key* for ``runs`` table.
136
+ """
137
+ name = Column(String(250))
138
+ """
139
+ Run name: `String` (limit 250 characters).
140
+ """
141
+ source_type = Column(String(20), default=SourceType.to_string(SourceType.LOCAL))
142
+ """
143
+ Source Type: `String` (limit 20 characters). Can be one of ``NOTEBOOK``, ``JOB``, ``PROJECT``,
144
+ ``LOCAL`` (default), or ``UNKNOWN``.
145
+ """
146
+ source_name = Column(String(500))
147
+ """
148
+ Name of source recording the run: `String` (limit 500 characters).
149
+ """
150
+ entry_point_name = Column(String(50))
151
+ """
152
+ Entry-point name that launched the run run: `String` (limit 50 characters).
153
+ """
154
+ user_id = Column(String(256), nullable=True, default=None)
155
+ """
156
+ User ID: `String` (limit 256 characters). Defaults to ``null``.
157
+ """
158
+ status = Column(String(20), default=RunStatus.to_string(RunStatus.SCHEDULED))
159
+ """
160
+ Run Status: `String` (limit 20 characters). Can be one of ``RUNNING``, ``SCHEDULED`` (default),
161
+ ``FINISHED``, ``FAILED``.
162
+ """
163
+ start_time = Column(BigInteger, default=get_current_time_millis)
164
+ """
165
+ Run start time: `BigInteger`. Defaults to current system time.
166
+ """
167
+ end_time = Column(BigInteger, nullable=True, default=None)
168
+ """
169
+ Run end time: `BigInteger`.
170
+ """
171
+ deleted_time = Column(BigInteger, nullable=True, default=None)
172
+ """
173
+ Run deleted time: `BigInteger`. Timestamp of when run is deleted, defaults to none.
174
+ """
175
+ source_version = Column(String(50))
176
+ """
177
+ Source version: `String` (limit 50 characters).
178
+ """
179
+ lifecycle_stage = Column(String(20), default=LifecycleStage.ACTIVE)
180
+ """
181
+ Lifecycle Stage of run: `String` (limit 32 characters).
182
+ Can be either ``active`` (default) or ``deleted``.
183
+ """
184
+ artifact_uri = Column(String(200), default=None)
185
+ """
186
+ Default artifact location for this run: `String` (limit 200 characters).
187
+ """
188
+ experiment_id = Column(Integer, ForeignKey("experiments.experiment_id"))
189
+ """
190
+ Experiment ID to which this run belongs to: *Foreign Key* into ``experiment`` table.
191
+ """
192
+ experiment = relationship("SqlExperiment", backref=backref("runs", cascade="all"))
193
+ """
194
+ SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlExperiment`.
195
+ """
196
+
197
+ __table_args__ = (
198
+ CheckConstraint(source_type.in_(SourceTypes), name="source_type"),
199
+ CheckConstraint(status.in_(RunStatusTypes), name="status"),
200
+ CheckConstraint(
201
+ lifecycle_stage.in_(LifecycleStage.view_type_to_stages(ViewType.ALL)),
202
+ name="runs_lifecycle_stage",
203
+ ),
204
+ PrimaryKeyConstraint("run_uuid", name="run_pk"),
205
+ )
206
+
207
+ @staticmethod
208
+ def get_attribute_name(mlflow_attribute_name):
209
+ """
210
+ Resolves an MLflow attribute name to a `SqlRun` attribute name.
211
+ """
212
+ # Currently, MLflow Search attributes defined in `SearchUtils.VALID_SEARCH_ATTRIBUTE_KEYS`
213
+ # share the same names as their corresponding `SqlRun` attributes. Therefore, this function
214
+ # returns the same attribute name
215
+ return {"run_name": "name", "run_id": "run_uuid"}.get(
216
+ mlflow_attribute_name, mlflow_attribute_name
217
+ )
218
+
219
+ def to_mlflow_entity(self):
220
+ """
221
+ Convert DB model to corresponding MLflow entity.
222
+
223
+ Returns:
224
+ mlflow.entities.Run: Description of the return value.
225
+ """
226
+ run_info = RunInfo(
227
+ run_id=self.run_uuid,
228
+ run_name=self.name,
229
+ experiment_id=str(self.experiment_id),
230
+ user_id=self.user_id,
231
+ status=self.status,
232
+ start_time=self.start_time,
233
+ end_time=self.end_time,
234
+ lifecycle_stage=self.lifecycle_stage,
235
+ artifact_uri=self.artifact_uri,
236
+ )
237
+
238
+ tags = [t.to_mlflow_entity() for t in self.tags]
239
+ run_data = RunData(
240
+ metrics=[m.to_mlflow_entity() for m in self.latest_metrics],
241
+ params=[p.to_mlflow_entity() for p in self.params],
242
+ tags=tags,
243
+ )
244
+ if not run_info.run_name:
245
+ run_name = _get_run_name_from_tags(tags)
246
+ if run_name:
247
+ run_info._set_run_name(run_name)
248
+
249
+ return Run(run_info=run_info, run_data=run_data)
250
+
251
+
252
+ class SqlExperimentTag(Base):
253
+ """
254
+ DB model for :py:class:`mlflow.entities.RunTag`.
255
+ These are recorded in ``experiment_tags`` table.
256
+ """
257
+
258
+ __tablename__ = "experiment_tags"
259
+
260
+ key = Column(String(250))
261
+ """
262
+ Tag key: `String` (limit 250 characters). *Primary Key* for ``tags`` table.
263
+ """
264
+ value = Column(String(5000), nullable=True)
265
+ """
266
+ Value associated with tag: `String` (limit 5000 characters). Could be *null*.
267
+ """
268
+ experiment_id = Column(Integer, ForeignKey("experiments.experiment_id"))
269
+ """
270
+ Experiment ID to which this tag belongs: *Foreign Key* into ``experiments`` table.
271
+ """
272
+ experiment = relationship("SqlExperiment", backref=backref("tags", cascade="all"))
273
+ """
274
+ SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlExperiment`.
275
+ """
276
+
277
+ __table_args__ = (PrimaryKeyConstraint("key", "experiment_id", name="experiment_tag_pk"),)
278
+
279
+ def __repr__(self):
280
+ return f"<SqlExperimentTag({self.key}, {self.value})>"
281
+
282
+ def to_mlflow_entity(self):
283
+ """
284
+ Convert DB model to corresponding MLflow entity.
285
+
286
+ Returns:
287
+ mlflow.entities.RunTag: Description of the return value.
288
+ """
289
+ return ExperimentTag(key=self.key, value=self.value)
290
+
291
+
292
+ class SqlTag(Base):
293
+ """
294
+ DB model for :py:class:`mlflow.entities.RunTag`. These are recorded in ``tags`` table.
295
+ """
296
+
297
+ __tablename__ = "tags"
298
+ __table_args__ = (
299
+ PrimaryKeyConstraint("key", "run_uuid", name="tag_pk"),
300
+ Index(f"index_{__tablename__}_run_uuid", "run_uuid"),
301
+ )
302
+
303
+ key = Column(String(250))
304
+ """
305
+ Tag key: `String` (limit 250 characters). *Primary Key* for ``tags`` table.
306
+ """
307
+ value = Column(String(8000), nullable=True)
308
+ """
309
+ Value associated with tag: `String` (limit 8000 characters). Could be *null*.
310
+ """
311
+ run_uuid = Column(String(32), ForeignKey("runs.run_uuid"))
312
+ """
313
+ Run UUID to which this tag belongs to: *Foreign Key* into ``runs`` table.
314
+ """
315
+ run = relationship("SqlRun", backref=backref("tags", cascade="all"))
316
+ """
317
+ SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlRun`.
318
+ """
319
+
320
+ def __repr__(self):
321
+ return f"<SqlRunTag({self.key}, {self.value})>"
322
+
323
+ def to_mlflow_entity(self):
324
+ """
325
+ Convert DB model to corresponding MLflow entity.
326
+
327
+ Returns:
328
+ :py:class:`mlflow.entities.RunTag`.
329
+ """
330
+ return RunTag(key=self.key, value=self.value)
331
+
332
+
333
+ class SqlMetric(Base):
334
+ __tablename__ = "metrics"
335
+ __table_args__ = (
336
+ PrimaryKeyConstraint(
337
+ "key", "timestamp", "step", "run_uuid", "value", "is_nan", name="metric_pk"
338
+ ),
339
+ Index(f"index_{__tablename__}_run_uuid", "run_uuid"),
340
+ )
341
+
342
+ key = Column(String(250))
343
+ """
344
+ Metric key: `String` (limit 250 characters). Part of *Primary Key* for ``metrics`` table.
345
+ """
346
+ value = Column(sa.types.Float(precision=53), nullable=False)
347
+ """
348
+ Metric value: `Float`. Defined as *Non-null* in schema.
349
+ """
350
+ timestamp = Column(BigInteger, default=get_current_time_millis)
351
+ """
352
+ Timestamp recorded for this metric entry: `BigInteger`. Part of *Primary Key* for
353
+ ``metrics`` table.
354
+ """
355
+ step = Column(BigInteger, default=0, nullable=False)
356
+ """
357
+ Step recorded for this metric entry: `BigInteger`.
358
+ """
359
+ is_nan = Column(Boolean(create_constraint=True), nullable=False, default=False)
360
+ """
361
+ True if the value is in fact NaN.
362
+ """
363
+ run_uuid = Column(String(32), ForeignKey("runs.run_uuid"))
364
+ """
365
+ Run UUID to which this metric belongs to: Part of *Primary Key* for ``metrics`` table.
366
+ *Foreign Key* into ``runs`` table.
367
+ """
368
+ run = relationship("SqlRun", backref=backref("metrics", cascade="all"))
369
+ """
370
+ SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlRun`.
371
+ """
372
+
373
+ def __repr__(self):
374
+ return f"<SqlMetric({self.key}, {self.value}, {self.timestamp}, {self.step})>"
375
+
376
+ def to_mlflow_entity(self):
377
+ """
378
+ Convert DB model to corresponding MLflow entity.
379
+
380
+ Returns:
381
+ mlflow.entities.Metric: Description of the return value.
382
+ """
383
+ return Metric(
384
+ key=self.key,
385
+ value=self.value if not self.is_nan else float("nan"),
386
+ timestamp=self.timestamp,
387
+ step=self.step,
388
+ )
389
+
390
+
391
+ class SqlLatestMetric(Base):
392
+ __tablename__ = "latest_metrics"
393
+ __table_args__ = (
394
+ PrimaryKeyConstraint("key", "run_uuid", name="latest_metric_pk"),
395
+ Index(f"index_{__tablename__}_run_uuid", "run_uuid"),
396
+ )
397
+
398
+ key = Column(String(250))
399
+ """
400
+ Metric key: `String` (limit 250 characters). Part of *Primary Key* for ``latest_metrics`` table.
401
+ """
402
+ value = Column(sa.types.Float(precision=53), nullable=False)
403
+ """
404
+ Metric value: `Float`. Defined as *Non-null* in schema.
405
+ """
406
+ timestamp = Column(BigInteger, default=get_current_time_millis)
407
+ """
408
+ Timestamp recorded for this metric entry: `BigInteger`. Part of *Primary Key* for
409
+ ``latest_metrics`` table.
410
+ """
411
+ step = Column(BigInteger, default=0, nullable=False)
412
+ """
413
+ Step recorded for this metric entry: `BigInteger`.
414
+ """
415
+ is_nan = Column(Boolean(create_constraint=True), nullable=False, default=False)
416
+ """
417
+ True if the value is in fact NaN.
418
+ """
419
+ run_uuid = Column(String(32), ForeignKey("runs.run_uuid"))
420
+ """
421
+ Run UUID to which this metric belongs to: Part of *Primary Key* for ``latest_metrics`` table.
422
+ *Foreign Key* into ``runs`` table.
423
+ """
424
+ run = relationship("SqlRun", backref=backref("latest_metrics", cascade="all"))
425
+ """
426
+ SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlRun`.
427
+ """
428
+
429
+ def __repr__(self):
430
+ return f"<SqlLatestMetric({self.key}, {self.value}, {self.timestamp}, {self.step})>"
431
+
432
+ def to_mlflow_entity(self):
433
+ """
434
+ Convert DB model to corresponding MLflow entity.
435
+
436
+ Returns:
437
+ mlflow.entities.Metric: Description of the return value.
438
+ """
439
+ return Metric(
440
+ key=self.key,
441
+ value=self.value if not self.is_nan else float("nan"),
442
+ timestamp=self.timestamp,
443
+ step=self.step,
444
+ )
445
+
446
+
447
+ class SqlParam(Base):
448
+ __tablename__ = "params"
449
+ __table_args__ = (
450
+ PrimaryKeyConstraint("key", "run_uuid", name="param_pk"),
451
+ Index(f"index_{__tablename__}_run_uuid", "run_uuid"),
452
+ )
453
+
454
+ key = Column(String(250))
455
+ """
456
+ Param key: `String` (limit 250 characters). Part of *Primary Key* for ``params`` table.
457
+ """
458
+ value = Column(String(8000), nullable=False)
459
+ """
460
+ Param value: `String` (limit 8000 characters). Defined as *Non-null* in schema.
461
+ """
462
+ run_uuid = Column(String(32), ForeignKey("runs.run_uuid"))
463
+ """
464
+ Run UUID to which this metric belongs to: Part of *Primary Key* for ``params`` table.
465
+ *Foreign Key* into ``runs`` table.
466
+ """
467
+ run = relationship("SqlRun", backref=backref("params", cascade="all"))
468
+ """
469
+ SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlRun`.
470
+ """
471
+
472
+ def __repr__(self):
473
+ return f"<SqlParam({self.key}, {self.value})>"
474
+
475
+ def to_mlflow_entity(self):
476
+ """
477
+ Convert DB model to corresponding MLflow entity.
478
+
479
+ Returns:
480
+ mlflow.entities.Param: Description of the return value.
481
+ """
482
+ return Param(key=self.key, value=self.value)
483
+
484
+
485
+ class SqlDataset(Base):
486
+ __tablename__ = "datasets"
487
+ __table_args__ = (
488
+ PrimaryKeyConstraint("experiment_id", "name", "digest", name="dataset_pk"),
489
+ Index(f"index_{__tablename__}_dataset_uuid", "dataset_uuid"),
490
+ Index(
491
+ f"index_{__tablename__}_experiment_id_dataset_source_type",
492
+ "experiment_id",
493
+ "dataset_source_type",
494
+ ),
495
+ )
496
+
497
+ dataset_uuid = Column(String(36), nullable=False)
498
+ """
499
+ Dataset UUID: `String` (limit 36 characters). Defined as *Non-null* in schema.
500
+ Part of *Primary Key* for ``datasets`` table.
501
+ """
502
+ experiment_id = Column(Integer, ForeignKey("experiments.experiment_id", ondelete="CASCADE"))
503
+ """
504
+ Experiment ID to which this dataset belongs: *Foreign Key* into ``experiments`` table.
505
+ """
506
+ name = Column(String(500), nullable=False)
507
+ """
508
+ Param name: `String` (limit 500 characters). Defined as *Non-null* in schema.
509
+ Part of *Primary Key* for ``datasets`` table.
510
+ """
511
+ digest = Column(String(36), nullable=False)
512
+ """
513
+ Param digest: `String` (limit 500 characters). Defined as *Non-null* in schema.
514
+ Part of *Primary Key* for ``datasets`` table.
515
+ """
516
+ dataset_source_type = Column(String(36), nullable=False)
517
+ """
518
+ Param dataset_source_type: `String` (limit 36 characters). Defined as *Non-null* in schema.
519
+ """
520
+ dataset_source = Column(UnicodeText, nullable=False)
521
+ """
522
+ Param dataset_source: `UnicodeText`. Defined as *Non-null* in schema.
523
+ """
524
+ dataset_schema = Column(UnicodeText, nullable=True)
525
+ """
526
+ Param dataset_schema: `UnicodeText`.
527
+ """
528
+ dataset_profile = Column(UnicodeText, nullable=True)
529
+ """
530
+ Param dataset_profile: `UnicodeText`.
531
+ """
532
+
533
+ def __repr__(self):
534
+ return "<SqlDataset ({}, {}, {}, {}, {}, {}, {}, {})>".format(
535
+ self.dataset_uuid,
536
+ self.experiment_id,
537
+ self.name,
538
+ self.digest,
539
+ self.dataset_source_type,
540
+ self.dataset_source,
541
+ self.dataset_schema,
542
+ self.dataset_profile,
543
+ )
544
+
545
+ def to_mlflow_entity(self):
546
+ """
547
+ Convert DB model to corresponding MLflow entity.
548
+
549
+ Returns:
550
+ mlflow.entities.Dataset.
551
+ """
552
+ return Dataset(
553
+ name=self.name,
554
+ digest=self.digest,
555
+ source_type=self.dataset_source_type,
556
+ source=self.dataset_source,
557
+ schema=self.dataset_schema,
558
+ profile=self.dataset_profile,
559
+ )
560
+
561
+
562
+ class SqlInput(Base):
563
+ __tablename__ = "inputs"
564
+ __table_args__ = (
565
+ PrimaryKeyConstraint(
566
+ "source_type", "source_id", "destination_type", "destination_id", name="inputs_pk"
567
+ ),
568
+ Index(f"index_{__tablename__}_input_uuid", "input_uuid"),
569
+ Index(
570
+ f"index_{__tablename__}_destination_type_destination_id_source_type",
571
+ "destination_type",
572
+ "destination_id",
573
+ "source_type",
574
+ ),
575
+ )
576
+
577
+ input_uuid = Column(String(36), nullable=False)
578
+ """
579
+ Input UUID: `String` (limit 36 characters). Defined as *Non-null* in schema.
580
+ """
581
+ source_type = Column(String(36), nullable=False)
582
+ """
583
+ Source type: `String` (limit 36 characters). Defined as *Non-null* in schema.
584
+ Part of *Primary Key* for ``inputs`` table.
585
+ """
586
+ source_id = Column(String(36), nullable=False)
587
+ """
588
+ Source Id: `String` (limit 36 characters). Defined as *Non-null* in schema.
589
+ Part of *Primary Key* for ``inputs`` table.
590
+ """
591
+ destination_type = Column(String(36), nullable=False)
592
+ """
593
+ Destination type: `String` (limit 36 characters). Defined as *Non-null* in schema.
594
+ Part of *Primary Key* for ``inputs`` table.
595
+ """
596
+ destination_id = Column(String(36), nullable=False)
597
+ """
598
+ Destination Id: `String` (limit 36 characters). Defined as *Non-null* in schema.
599
+ Part of *Primary Key* for ``inputs`` table.
600
+ """
601
+ step = Column(BigInteger, nullable=False, server_default="0")
602
+
603
+ def __repr__(self):
604
+ return "<SqlInput ({}, {}, {}, {}, {})>".format(
605
+ self.input_uuid,
606
+ self.source_type,
607
+ self.source_id,
608
+ self.destination_type,
609
+ self.destination_id,
610
+ )
611
+
612
+
613
+ class SqlInputTag(Base):
614
+ __tablename__ = "input_tags"
615
+ __table_args__ = (PrimaryKeyConstraint("input_uuid", "name", name="input_tags_pk"),)
616
+
617
+ input_uuid = Column(String(36), ForeignKey("inputs.input_uuid"), nullable=False)
618
+ """
619
+ Input UUID: `String` (limit 36 characters). Defined as *Non-null* in schema.
620
+ *Foreign Key* into ``inputs`` table. Part of *Primary Key* for ``input_tags`` table.
621
+ """
622
+ name = Column(String(255), nullable=False)
623
+ """
624
+ Param name: `String` (limit 255 characters). Defined as *Non-null* in schema.
625
+ Part of *Primary Key* for ``input_tags`` table.
626
+ """
627
+ value = Column(String(500), nullable=False)
628
+ """
629
+ Param value: `String` (limit 500 characters). Defined as *Non-null* in schema.
630
+ Part of *Primary Key* for ``input_tags`` table.
631
+ """
632
+
633
+ def __repr__(self):
634
+ return f"<SqlInputTag ({self.input_uuid}, {self.name}, {self.value})>"
635
+
636
+ def to_mlflow_entity(self):
637
+ """
638
+ Convert DB model to corresponding MLflow entity.
639
+
640
+ Returns:
641
+ mlflow.entities.InputTag: Description of the return value.
642
+ """
643
+ return InputTag(key=self.name, value=self.value)
644
+
645
+
646
+ #######################################################################################
647
+ # Below are Tracing models. We may refactor them to be in a separate module in the future.
648
+ #######################################################################################
649
+
650
+
651
+ class SqlTraceInfo(Base):
652
+ __tablename__ = "trace_info"
653
+
654
+ request_id = Column(String(50), nullable=False)
655
+ """
656
+ Trace ID: `String` (limit 50 characters). *Primary Key* for ``trace_info`` table.
657
+ Named as "trace_id" in V3 format.
658
+ """
659
+ experiment_id = Column(Integer, ForeignKey("experiments.experiment_id"), nullable=False)
660
+ """
661
+ Experiment ID to which this trace belongs: *Foreign Key* into ``experiments`` table.
662
+ """
663
+ timestamp_ms = Column(BigInteger, nullable=False)
664
+ """
665
+ Start time of the trace, in milliseconds. Named as "request_time" in V3 format.
666
+ """
667
+ execution_time_ms = Column(BigInteger, nullable=True)
668
+ """
669
+ Duration of the trace, in milliseconds. Could be *null* if the trace is still in progress
670
+ or not ended correctly for some reason. Named as "execution_duration" in V3 format.
671
+ """
672
+ status = Column(String(50), nullable=False)
673
+ """
674
+ State of the trace. The values are defined in
675
+ :py:class:`mlflow.entities.trace_status.TraceStatus` enum but we don't enforce
676
+ constraint at DB level. Named as "state" in V3 format.
677
+ """
678
+ client_request_id = Column(String(50), nullable=True)
679
+ """
680
+ Client request ID: `String` (limit 50 characters). Could be *null*. Newly added in V3 format.
681
+ """
682
+ request_preview = Column(String(1000), nullable=True)
683
+ """
684
+ Request preview: `String` (limit 1000 characters). Could be *null*. Newly added in V3 format.
685
+ """
686
+ response_preview = Column(String(1000), nullable=True)
687
+ """
688
+ Response preview: `String` (limit 1000 characters). Could be *null*. Newly added in V3 format.
689
+ """
690
+
691
+ __table_args__ = (
692
+ PrimaryKeyConstraint("request_id", name="trace_info_pk"),
693
+ # The most frequent query will be get all traces in an experiment sorted by timestamp desc,
694
+ # which is the default view in the UI. Also every search query should have experiment_id(s)
695
+ # in the where clause.
696
+ Index(f"index_{__tablename__}_experiment_id_timestamp_ms", "experiment_id", "timestamp_ms"),
697
+ )
698
+
699
+ def to_mlflow_entity(self):
700
+ """
701
+ Convert DB model to corresponding MLflow entity.
702
+
703
+ Returns:
704
+ :py:class:`mlflow.entities.TraceInfo` object.
705
+ """
706
+ return TraceInfo(
707
+ trace_id=self.request_id,
708
+ trace_location=TraceLocation.from_experiment_id(str(self.experiment_id)),
709
+ request_time=self.timestamp_ms,
710
+ execution_duration=self.execution_time_ms,
711
+ state=TraceState(self.status),
712
+ tags={t.key: t.value for t in self.tags},
713
+ trace_metadata={m.key: m.value for m in self.request_metadata},
714
+ client_request_id=self.client_request_id,
715
+ request_preview=self.request_preview,
716
+ response_preview=self.response_preview,
717
+ assessments=[], # Implement this once we support assessments in OSS
718
+ )
719
+
720
+
721
+ class SqlTraceTag(Base):
722
+ __tablename__ = "trace_tags"
723
+
724
+ key = Column(String(250))
725
+ """
726
+ Tag key: `String` (limit 250 characters).
727
+ """
728
+ value = Column(String(8000), nullable=True)
729
+ """
730
+ Value associated with tag: `String` (limit 250 characters). Could be *null*.
731
+ """
732
+ request_id = Column(
733
+ String(50), ForeignKey("trace_info.request_id", ondelete="CASCADE"), nullable=False
734
+ )
735
+ """
736
+ Request ID to which this tag belongs: *Foreign Key* into ``trace_info`` table.
737
+ """
738
+ trace_info = relationship("SqlTraceInfo", backref=backref("tags", cascade="all"))
739
+ """
740
+ SQLAlchemy relationship (many:one) with
741
+ :py:class:`mlflow.store.dbmodels.models.SqlTraceInfo`.
742
+ """
743
+
744
+ # Key is unique within a request_id
745
+ __table_args__ = (
746
+ PrimaryKeyConstraint("request_id", "key", name="trace_tag_pk"),
747
+ Index(f"index_{__tablename__}_request_id"),
748
+ )
749
+
750
+
751
+ class SqlTraceMetadata(Base):
752
+ __tablename__ = "trace_request_metadata"
753
+
754
+ key = Column(String(250))
755
+ """
756
+ Metadata key: `String` (limit 250 characters).
757
+ """
758
+ value = Column(String(8000), nullable=True)
759
+ """
760
+ Value associated with metadata: `String` (limit 250 characters). Could be *null*.
761
+ """
762
+ request_id = Column(
763
+ String(50), ForeignKey("trace_info.request_id", ondelete="CASCADE"), nullable=False
764
+ )
765
+ """
766
+ Request ID to which this metadata belongs: *Foreign Key* into ``trace_info`` table.
767
+ **Corresponding to the "trace_id" in V3 format.**
768
+ """
769
+ trace_info = relationship("SqlTraceInfo", backref=backref("request_metadata", cascade="all"))
770
+ """
771
+ SQLAlchemy relationship (many:one) with
772
+ :py:class:`mlflow.store.dbmodels.models.SqlTraceInfo`.
773
+ """
774
+
775
+ # Key is unique within a request_id
776
+ __table_args__ = (
777
+ PrimaryKeyConstraint("request_id", "key", name="trace_request_metadata_pk"),
778
+ Index(f"index_{__tablename__}_request_id"),
779
+ )
780
+
781
+
782
+ class SqlLoggedModel(Base):
783
+ __tablename__ = "logged_models"
784
+
785
+ model_id = Column(String(36), nullable=False)
786
+ """
787
+ Model ID: `String` (limit 36 characters). *Primary Key* for ``logged_models`` table.
788
+ """
789
+
790
+ experiment_id = Column(Integer, nullable=False)
791
+ """
792
+ Experiment ID to which this model belongs: *Foreign Key* into ``experiments`` table.
793
+ """
794
+
795
+ name = Column(String(500), nullable=False)
796
+ """
797
+ Model name: `String` (limit 500 characters).
798
+ """
799
+
800
+ artifact_location = Column(String(1000), nullable=False)
801
+ """
802
+ Artifact location: `String` (limit 1000 characters).
803
+ """
804
+
805
+ creation_timestamp_ms = Column(BigInteger, nullable=False)
806
+ """
807
+ Creation timestamp: `BigInteger`.
808
+ """
809
+
810
+ last_updated_timestamp_ms = Column(BigInteger, nullable=False)
811
+ """
812
+ Last updated timestamp: `BigInteger`.
813
+ """
814
+
815
+ status = Column(Integer, nullable=False)
816
+ """
817
+ Status: `Integer`.
818
+ """
819
+
820
+ lifecycle_stage = Column(String(32), default=LifecycleStage.ACTIVE)
821
+ """
822
+ Lifecycle Stage of model: `String` (limit 32 characters).
823
+ """
824
+
825
+ model_type = Column(String(500), nullable=True)
826
+ """
827
+ Model type: `String` (limit 500 characters).
828
+ """
829
+
830
+ source_run_id = Column(String(32), nullable=True)
831
+ """
832
+ Source run ID: `String` (limit 32 characters).
833
+ """
834
+
835
+ status_message = Column(String(1000), nullable=True)
836
+ """
837
+ Status message: `String` (limit 1000 characters).
838
+ """
839
+
840
+ tags = relationship("SqlLoggedModelTag", backref="logged_model", cascade="all")
841
+ params = relationship("SqlLoggedModelParam", backref="logged_model", cascade="all")
842
+ metrics = relationship("SqlLoggedModelMetric", backref="logged_model", cascade="all")
843
+
844
+ __table_args__ = (
845
+ PrimaryKeyConstraint("model_id", name="logged_models_pk"),
846
+ CheckConstraint(
847
+ lifecycle_stage.in_(LifecycleStage.view_type_to_stages(ViewType.ALL)),
848
+ name="logged_models_lifecycle_stage_check",
849
+ ),
850
+ ForeignKeyConstraint(
851
+ ["experiment_id"],
852
+ ["experiments.experiment_id"],
853
+ ondelete="CASCADE",
854
+ name="fk_logged_models_experiment_id",
855
+ ),
856
+ )
857
+
858
+ def to_mlflow_entity(self) -> LoggedModel:
859
+ return LoggedModel(
860
+ model_id=self.model_id,
861
+ experiment_id=str(self.experiment_id),
862
+ name=self.name,
863
+ artifact_location=self.artifact_location,
864
+ creation_timestamp=self.creation_timestamp_ms,
865
+ last_updated_timestamp=self.last_updated_timestamp_ms,
866
+ status=LoggedModelStatus.from_int(self.status),
867
+ model_type=self.model_type,
868
+ source_run_id=self.source_run_id,
869
+ status_message=self.status_message,
870
+ tags={t.tag_key: t.tag_value for t in self.tags} if self.tags else None,
871
+ params={p.param_key: p.param_value for p in self.params} if self.params else None,
872
+ metrics=[m.to_mlflow_entity() for m in self.metrics] if self.metrics else None,
873
+ )
874
+
875
+ ALIASES = {
876
+ "creation_time": "creation_timestamp_ms",
877
+ "creation_timestamp": "creation_timestamp_ms",
878
+ "last_updated_timestamp": "last_updated_timestamp_ms",
879
+ }
880
+
881
+ @staticmethod
882
+ def is_numeric(s: str) -> bool:
883
+ return SqlLoggedModel.ALIASES.get(s, s) in {
884
+ "creation_timestamp_ms",
885
+ "last_updated_timestamp_ms",
886
+ }
887
+
888
+
889
+ class SqlLoggedModelMetric(Base):
890
+ __tablename__ = "logged_model_metrics"
891
+
892
+ model_id = Column(String(36), nullable=False)
893
+ """
894
+ Model ID: `String` (limit 36 characters).
895
+ """
896
+
897
+ metric_name = Column(String(500), nullable=False)
898
+ """
899
+ Metric name: `String` (limit 500 characters).
900
+ """
901
+
902
+ metric_timestamp_ms = Column(BigInteger, nullable=False)
903
+ """
904
+ Metric timestamp: `BigInteger`.
905
+ """
906
+
907
+ metric_step = Column(BigInteger, nullable=False)
908
+ """
909
+ Metric step: `BigInteger`.
910
+ """
911
+
912
+ metric_value = Column(sa.types.Float(precision=53), nullable=True)
913
+ """
914
+ Metric value: `Float`.
915
+ """
916
+
917
+ experiment_id = Column(Integer, nullable=False)
918
+ """
919
+ Experiment ID: `Integer`.
920
+ """
921
+
922
+ run_id = Column(String(32), nullable=False)
923
+ """
924
+ Run ID: `String` (limit 32 characters).
925
+ """
926
+
927
+ dataset_uuid = Column(String(36), nullable=True)
928
+ """
929
+ Dataset UUID: `String` (limit 36 characters).
930
+ """
931
+
932
+ dataset_name = Column(String(500), nullable=True)
933
+ """
934
+ Dataset name: `String` (limit 500 characters).
935
+ """
936
+
937
+ dataset_digest = Column(String(36), nullable=True)
938
+ """
939
+ Dataset digest: `String` (limit 36 characters).
940
+ """
941
+
942
+ __table_args__ = (
943
+ PrimaryKeyConstraint(
944
+ "model_id",
945
+ "metric_name",
946
+ "metric_timestamp_ms",
947
+ "metric_step",
948
+ "run_id",
949
+ name="logged_model_metrics_pk",
950
+ ),
951
+ ForeignKeyConstraint(
952
+ ["model_id"],
953
+ ["logged_models.model_id"],
954
+ ondelete="CASCADE",
955
+ name="fk_logged_model_metrics_model_id",
956
+ ),
957
+ ForeignKeyConstraint(
958
+ ["experiment_id"],
959
+ ["experiments.experiment_id"],
960
+ name="fk_logged_model_metrics_experiment_id",
961
+ ),
962
+ ForeignKeyConstraint(
963
+ ["run_id"],
964
+ ["runs.run_uuid"],
965
+ ondelete="CASCADE",
966
+ name="fk_logged_model_metrics_run_id",
967
+ ),
968
+ Index("index_logged_model_metrics_model_id", "model_id"),
969
+ )
970
+
971
+ def to_mlflow_entity(self) -> Metric:
972
+ return Metric(
973
+ key=self.metric_name,
974
+ value=self.metric_value,
975
+ timestamp=self.metric_timestamp_ms,
976
+ step=self.metric_step,
977
+ run_id=self.run_id,
978
+ dataset_name=self.dataset_name,
979
+ dataset_digest=self.dataset_digest,
980
+ model_id=self.model_id,
981
+ )
982
+
983
+
984
+ class SqlLoggedModelParam(Base):
985
+ __tablename__ = "logged_model_params"
986
+
987
+ model_id = Column(String(36), nullable=False)
988
+ """
989
+ Model ID: `String` (limit 36 characters).
990
+ """
991
+
992
+ experiment_id = Column(Integer, nullable=False)
993
+ """
994
+ Experiment ID: `Integer`.
995
+ """
996
+
997
+ param_key = Column(String(255), nullable=False)
998
+ """
999
+ Param key: `String` (limit 255 characters).
1000
+ """
1001
+
1002
+ param_value = Column(Text(), nullable=False)
1003
+ """
1004
+ Param value: `Text`.
1005
+ """
1006
+
1007
+ __table_args__ = (
1008
+ PrimaryKeyConstraint(
1009
+ "model_id",
1010
+ "param_key",
1011
+ name="logged_model_params_pk",
1012
+ ),
1013
+ ForeignKeyConstraint(
1014
+ ["model_id"],
1015
+ ["logged_models.model_id"],
1016
+ name="fk_logged_model_params_model_id",
1017
+ ondelete="CASCADE",
1018
+ ),
1019
+ ForeignKeyConstraint(
1020
+ ["experiment_id"],
1021
+ ["experiments.experiment_id"],
1022
+ name="fk_logged_model_params_experiment_id",
1023
+ ),
1024
+ )
1025
+
1026
+ def to_mlflow_entity(self) -> LoggedModelParameter:
1027
+ return LoggedModelParameter(key=self.param_key, value=self.param_value)
1028
+
1029
+
1030
+ class SqlLoggedModelTag(Base):
1031
+ __tablename__ = "logged_model_tags"
1032
+
1033
+ model_id = Column(String(36), nullable=False)
1034
+ """
1035
+ Model ID: `String` (limit 36 characters).
1036
+ """
1037
+
1038
+ experiment_id = Column(Integer, nullable=False)
1039
+ """
1040
+ Experiment ID: `Integer`.
1041
+ """
1042
+
1043
+ tag_key = Column(String(255), nullable=False)
1044
+ """
1045
+ Tag key: `String` (limit 255 characters).
1046
+ """
1047
+
1048
+ tag_value = Column(Text(), nullable=False)
1049
+ """
1050
+ Tag value: `Text`.
1051
+ """
1052
+
1053
+ __table_args__ = (
1054
+ PrimaryKeyConstraint(
1055
+ "model_id",
1056
+ "tag_key",
1057
+ name="logged_model_tags_pk",
1058
+ ),
1059
+ ForeignKeyConstraint(
1060
+ ["model_id"],
1061
+ ["logged_models.model_id"],
1062
+ name="fk_logged_model_tags_model_id",
1063
+ ondelete="CASCADE",
1064
+ ),
1065
+ ForeignKeyConstraint(
1066
+ ["experiment_id"],
1067
+ ["experiments.experiment_id"],
1068
+ name="fk_logged_model_tags_experiment_id",
1069
+ ),
1070
+ )
1071
+
1072
+ def to_mlflow_entity(self) -> LoggedModelTag:
1073
+ return LoggedModelTag(key=self.tag_key, value=self.tag_value)