SUBMARINE-1023. Submarine SDK sqlalchemy store (model registry)

### What is this PR for?
* Implement the model registry SQL method in Python SDK
* Apply sqlalchemy mypy checks
* Replace submarine tracking_uri with db_uri

### What type of PR is it?
[Feature]

### Todos

### What is the Jira issue?
https://issues.apache.org/jira/browse/SUBMARINE-1023

### How should this be tested?
All of the tests are provided in `submarine-sdk/pysubmarine/tests/store/model_registry/test_sqlalchemy_store.py`

### Screenshots (if appropriate)

### Questions:
* Do the license files need updating? No
* Are there breaking changes for older versions? No
* Does this need new documentation? No

Author: KUAN-HSUN-LI <[email protected]>

Signed-off-by: Kevin <[email protected]>

Closes #752 from KUAN-HSUN-LI/SUBMARINE-1023 and squashes the following commits:

7723788 [KUAN-HSUN-LI] SUBMARINE-1023. Submarine model registry Python SDK
5e086b5 [KUAN-HSUN-LI] fix
8c6574d [KUAN-HSUN-LI] fix
dd41039 [KUAN-HSUN-LI] replace tacking_uri with db_uri
ad9c5d7 [KUAN-HSUN-LI] add init file
374aa2d [KUAN-HSUN-LI] SUBMARINE-1023. Submarine model registry Python SDK
b502fc1 [KUAN-HSUN-LI] SUBMARINE-1023. sqlalchemy store tests
343780f [KUAN-HSUN-LI] apply sqlalchmy mypy

dev-support/database/submarine-model.sql

@@ -31,26 +31,27 @@ CREATE TABLE `registered_model_tag` (
 
 DROP TABLE IF EXISTS `model_version`;
 CREATE TABLE `model_version` (
-	`name` VARCHAR(256) NOT NULL,
+	`name` VARCHAR(256) NOT NULL COMMENT 'Name of model',
 	`version` INTEGER NOT NULL,
+	`source` VARCHAR(512) NOT NULL COMMENT 'Model saved link',
 	`user_id` VARCHAR(64) NOT NULL COMMENT 'Id of the created user',
 	`experiment_id` VARCHAR(64) NOT NULL,
-	`current_stage` VARCHAR(20) COMMENT 'Model stage ex: None, production...',
+	`current_stage` VARCHAR(64) COMMENT 'Model stage ex: None, production...',
 	`creation_time` DATETIME(3) COMMENT 'Millisecond precision',
 	`last_updated_time` DATETIME(3) COMMENT 'Millisecond precision',
-	`source` VARCHAR(512) COMMENT 'Model saved link',
 	`dataset` VARCHAR(256) COMMENT 'Which dataset is used',
 	`description` VARCHAR(5000),
 	CONSTRAINT `model_version_pk` PRIMARY KEY (`name`, `version`),
-	FOREIGN KEY(`name`) REFERENCES `registered_model` (`name`) ON UPDATE CASCADE ON DELETE CASCADE
+	FOREIGN KEY(`name`) REFERENCES `registered_model` (`name`) ON UPDATE CASCADE ON DELETE CASCADE,
+	UNIQUE(`source`)
 ) ENGINE=InnoDB DEFAULT CHARSET=utf8;
 
-DROP TABLE IF EXISTS `model_tag`;
-CREATE TABLE `model_tag` (
-	`name` VARCHAR(256) NOT NULL,
+DROP TABLE IF EXISTS `model_version_tag`;
+CREATE TABLE `model_version_tag` (
+	`name` VARCHAR(256) NOT NULL COMMENT 'Name of model',
 	`version` INTEGER NOT NULL,
 	`tag` VARCHAR(256) NOT NULL,
-	CONSTRAINT `model_tag_pk` PRIMARY KEY (`name`, `version`, `tag`),
+	CONSTRAINT `model_version_tag_pk` PRIMARY KEY (`name`, `version`, `tag`),
 	FOREIGN KEY(`name`, `version`) REFERENCES `model_version` (`name`, `version`) ON UPDATE CASCADE ON DELETE CASCADE
 ) ENGINE=InnoDB DEFAULT CHARSET=utf8;
 
@@ -72,7 +73,7 @@ DROP TABLE IF EXISTS `param`;
 CREATE TABLE `param` (
 	`id` VARCHAR(64) NOT NULL COMMENT 'Id of the Experiment',
 	`key` VARCHAR(190) NOT NULL COMMENT '`String` (limit 190 characters). Part of *Primary Key* for ``param`` table.',
-	`value` VARCHAR(32) NOT NULL COMMENT '`String` (limit 190 characters). Defined as *Non-null* in schema.',
+	`value` VARCHAR(190) NOT NULL COMMENT '`String` (limit 190 characters). Defined as *Non-null* in schema.',
 	`worker_index` VARCHAR(32) NOT NULL COMMENT '`String` (limit 32 characters). Part of *Primary Key* for\r\n    ``metric`` table.',
 	CONSTRAINT `param_pk` PRIMARY KEY  (`id`, `key`, `worker_index`),
 	FOREIGN KEY(`id`) REFERENCES `experiment` (`id`) ON UPDATE CASCADE ON DELETE CASCADE

dev-support/style-check/python/mypy-requirements.txt

@@ -18,3 +18,4 @@ types-requests==2.25.6
 types-certifi==2020.4.0
 types-six==1.16.1
 types-python-dateutil==2.8.0
+sqlalchemy[mypy]

pyproject.toml

@@ -20,4 +20,6 @@ profile = "black"
 line_length = 100
 [tool.black]
 max-line-length = 100
-line-length = 100
+line-length = 100
+[tool.mypy]
+plugins = "sqlalchemy.ext.mypy.plugin"

submarine-sdk/pysubmarine/github-actions/test-requirements.txt

@@ -23,11 +23,11 @@ pytest==3.2.1
 pytest-cov==2.6.0
 pytest-localserver==0.5.0
 pylint==2.5.2
-sqlalchemy==1.3.0
 PyMySQL==0.9.3
 pytest-mock==1.13.0
 certifi >= 14.05.14
 six >= 1.10
 python_dateutil >= 2.5.3
 setuptools >= 21.0.0
 urllib3 >= 1.15.1
+freezegun==1.1.0

submarine-sdk/pysubmarine/setup.py

@@ -30,7 +30,7 @@
         "six>=1.10.0",
         "numpy==1.18.5",
         "pandas",
-        "sqlalchemy",
+        "sqlalchemy>=1.4.0",
         "sqlparse",
         "pymysql",
         "requests==2.26.0",

submarine-sdk/pysubmarine/submarine/__init__.py

@@ -13,21 +13,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import submarine.tracking as tracking
 import submarine.tracking.fluent
+import submarine.utils as utils
 from submarine.experiment.api.experiment_client import ExperimentClient
 from submarine.models.client import ModelsClient
 
 log_param = submarine.tracking.fluent.log_param
 log_metric = submarine.tracking.fluent.log_metric
-set_tracking_uri = tracking.set_tracking_uri
-get_tracking_uri = tracking.get_tracking_uri
+set_db_uri = utils.set_db_uri
+get_db_uri = utils.get_db_uri
 
 __all__ = [
     "log_metric",
     "log_param",
-    "set_tracking_uri",
-    "get_tracking_uri",
+    "set_db_uri",
+    "get_db_uri",
     "ExperimentClient",
     "ModelsClient",
 ]

submarine-sdk/pysubmarine/submarine/entities/model_registry/__init__.py

@@ -13,14 +13,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from submarine.entities.model_registry.model_tag import ModelTag
 from submarine.entities.model_registry.model_version import ModelVersion
+from submarine.entities.model_registry.model_version_tag import ModelVersionTag
 from submarine.entities.model_registry.registered_model import RegisteredModel
 from submarine.entities.model_registry.registered_model_tag import RegisteredModelTag
 
 __all__ = [
-    "ModelVersion",
-    "ModelTag",
     "RegisteredModel",
     "RegisteredModelTag",
+    "ModelVersion",
+    "ModelVersionTag",
 ]

submarine-sdk/pysubmarine/submarine/entities/model_registry/model_stages.py

@@ -0,0 +1,33 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from submarine.exceptions import SubmarineException
+
+STAGE_NONE = "None"
+STAGE_DEVELOPING = "Developing"
+STAGE_PRODUCTION = "Production"
+STAGE_ARCHIVED = "Archived"
+
+STAGE_DELETED_INTERNAL = "Deleted_Internal"
+
+ALL_STAGES = [STAGE_NONE, STAGE_DEVELOPING, STAGE_PRODUCTION, STAGE_ARCHIVED]
+_CANONICAL_MAPPING = {stage.lower(): stage for stage in ALL_STAGES}
+
+
+def get_canonical_stage(stage):
+    key = stage.lower()
+    if key not in _CANONICAL_MAPPING:
+        raise SubmarineException(f"Invalid Model Version stage {stage}.")
+    return _CANONICAL_MAPPING[key]

submarine-sdk/pysubmarine/submarine/entities/model_registry/model_version.py

@@ -18,78 +18,78 @@
 
 class ModelVersion(_SubmarineObject):
     """
-    Model Version object.
+    Model version object.
     """
 
     def __init__(
         self,
         name,
         version,
+        source,
         user_id,
         experiment_id,
         current_stage,
         creation_time,
         last_updated_time,
-        source,
         dataset=None,
         description=None,
         tags=None,
     ):
         self._name = name
         self._version = version
+        self._source = source
         self._user_id = user_id
         self._experiment_id = experiment_id
         self._current_stage = current_stage
         self._creation_time = creation_time
         self._last_updated_time = last_updated_time
-        self._source = source
         self._dataset = dataset
         self._description = description
         self._tags = [tag.tag for tag in (tags or [])]
 
     @property
     def name(self):
-        """String. Unique name within Model Registry."""
+        """String. Registered model name"""
         return self._name
 
     @property
     def version(self):
-        """String. version"""
+        """Integer. version"""
         return self._version
 
+    @property
+    def source(self):
+        """String. Source path for the model."""
+        return self._source
+
     @property
     def user_id(self):
-        """String. User ID that created this model version."""
+        """String. User ID that created this version."""
         return self._user_id
 
     @property
     def experiment_id(self):
-        """String. Experiment ID that created this model version."""
+        """String. Experiment ID that created this version."""
         return self._experiment_id
 
     @property
     def creation_time(self):
-        """Datetime object. Model version creation timestamp."""
+        """Datetime object. The creation datetime of this version."""
         return self._creation_time
 
     @property
     def last_updated_time(self):
-        """Datetime object. Timestamp of last update for this model version."""
+        """Datetime object. Datetime of last update for this version."""
         return self._last_updated_time
 
-    @property
-    def source(self):
-        """String. Source path for the model."""
-        return self._source
-
     @property
     def current_stage(self):
-        """String. Current stage of this model version."""
+        """String. Current stage of this version."""
         return self._current_stage
 
     @property
     def dataset(self):
-        """String. Dataset used by this model version"""
+        """String. Dataset used for this version."""
         return self._dataset
 
     @property
@@ -99,8 +99,5 @@ def description(self):
 
     @property
     def tags(self):
-        """List of strings"""
+        """List of strings."""
         return self._tags
-
-    def _add_tag(self, tag):
-        self._tags.append(tag)

submarine-sdk/pysubmarine/submarine/entities/model_registry/model_tag.py → submarine-sdk/pysubmarine/submarine/entities/model_registry/model_version_tag.py

@@ -16,7 +16,7 @@
 from submarine.entities._submarine_object import _SubmarineObject
 
 
-class ModelTag(_SubmarineObject):
+class ModelVersionTag(_SubmarineObject):
     """
     Tag object associated with a model version.
     """
@@ -26,5 +26,5 @@ def __init__(self, tag):
 
     @property
     def tag(self):
-        """String tag"""
+        """String tag."""
         return self._tag

submarine-sdk/pysubmarine/submarine/entities/model_registry/registered_model.py

@@ -18,7 +18,7 @@
 
 class RegisteredModel(_SubmarineObject):
     """
-    Registered Model object.
+    Registered model object.
     """
 
     def __init__(self, name, creation_time, last_updated_time, description=None, tags=None):
@@ -35,12 +35,12 @@ def name(self):
 
     @property
     def creation_time(self):
-        """Datetime object. Model version creation timestamp."""
+        """Datetime object. Registered model creation datetime."""
         return self._creation_time
 
     @property
     def last_updated_time(self):
-        """Datetime object. Timestamp of last update for this model version."""
+        """Datetime object. Datetime of last update for this model."""
         return self._last_updated_time
 
     @property
@@ -52,6 +52,3 @@ def description(self):
     def tags(self):
         """List of strings"""
         return self._tags
-
-    def _add_tag(self, tag):
-        self._tags.append(tag)

submarine-sdk/pysubmarine/submarine/entities/model_registry/registered_model_tag.py

@@ -26,5 +26,5 @@ def __init__(self, tag):
 
     @property
     def tag(self):
-        """String tag."""
+        """String tag"""
         return self._tag