SUBMARINE-1045. Add static type parameter in submarine-sdk

submarine-sdk/pysubmarine/submarine/artifacts/repository.py

@@ -19,7 +19,7 @@
 
 
 class Repository:
-    def __init__(self, experiment_id):
+    def __init__(self, experiment_id: str) -> None:
         self.client = boto3.client(
             "s3",
             aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID"),
@@ -28,7 +28,7 @@ def __init__(self, experiment_id):
         )
         self.dest_path = experiment_id
 
-    def _upload_file(self, local_file, bucket, key):
+    def _upload_file(self, local_file: str, bucket: str, key: str) -> None:
         self.client.upload_file(Filename=local_file, Bucket=bucket, Key=key)
 
     def _list_artifact_subfolder(self, artifact_path):
@@ -39,7 +39,7 @@ def _list_artifact_subfolder(self, artifact_path):
         )
         return response.get("CommonPrefixes")
 
-    def log_artifact(self, local_file, artifact_path):
+    def log_artifact(self, local_file: str, artifact_path: str) -> None:
         bucket = "submarine"
         dest_path = self.dest_path
         dest_path = os.path.join(dest_path, artifact_path)
@@ -50,7 +50,7 @@ def log_artifact(self, local_file, artifact_path):
             key=dest_path,
         )
 
-    def log_artifacts(self, local_dir, artifact_path):
+    def log_artifacts(self, local_dir: str, artifact_path: str) -> None:
         bucket = "submarine"
         dest_path = self.dest_path
         list_of_subfolder = self._list_artifact_subfolder(artifact_path)

submarine-sdk/pysubmarine/submarine/entities/Metric.py

@@ -13,6 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Type
+
+from sqlalchemy.sql.schema import Column
+
 from submarine.entities._submarine_object import _SubmarineObject
 
 
@@ -21,34 +25,41 @@ class Metric(_SubmarineObject):
     Metric object.
     """
 
-    def __init__(self, key, value, worker_index, timestamp, step):
+    def __init__(
+        self,
+        key: Type[Column],
+        value: Type[Column],
+        worker_index: Type[Column],
+        timestamp: Type[Column],
+        step: Type[Column],
+    ):
         self._key = key
         self._value = value
         self._worker_index = worker_index
         self._timestamp = timestamp
         self._step = step
 
     @property
-    def key(self):
+    def key(self) -> Type[Column]:
         """String key corresponding to the metric name."""
         return self._key
 
     @property
-    def value(self):
+    def value(self) -> Type[Column]:
         """Float value of the metric."""
         return self._value
 
     @property
-    def worker_index(self):
+    def worker_index(self) -> Type[Column]:
         """string value of the metric."""
         return self._worker_index
 
     @property
-    def timestamp(self):
+    def timestamp(self) -> Type[Column]:
         """Metric timestamp as aa datetime object."""
         return self._timestamp
 
     @property
-    def step(self):
+    def step(self) -> Type[Column]:
         """Integer metric step (x-coordinate)."""
         return self._step

submarine-sdk/pysubmarine/submarine/entities/Param.py

@@ -13,6 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Type
+
+from sqlalchemy.sql.schema import Column
+
 from submarine.entities._submarine_object import _SubmarineObject
 
 
@@ -21,22 +25,22 @@ class Param(_SubmarineObject):
     Parameter object.
     """
 
-    def __init__(self, key, value, worker_index):
+    def __init__(self, key: Type[Column], value: Type[Column], worker_index: Type[Column]):
         self._key = key
         self._value = value
         self._worker_index = worker_index
 
     @property
-    def key(self):
+    def key(self) -> Type[Column]:
         """String key corresponding to the parameter name."""
         return self._key
 
     @property
-    def value(self):
+    def value(self) -> Type[Column]:
         """String value of the parameter."""
         return self._value
 
     @property
-    def worker_index(self):
+    def worker_index(self) -> Type[Column]:
         """String value of the parameter."""
         return self._worker_index

submarine-sdk/pysubmarine/submarine/entities/_submarine_object.py

@@ -31,11 +31,11 @@ def from_dictionary(cls, the_dict):
         filtered_dict = {key: value for key, value in the_dict.items() if key in cls._properties()}
         return cls(**filtered_dict)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return to_string(self)
 
 
-def to_string(obj):
+def to_string(obj) -> str:
     return _SubmarineObjectPrinter().to_string(obj)
 
 
@@ -48,10 +48,10 @@ def __init__(self):
         super(_SubmarineObjectPrinter, self).__init__()
         self.printer = pprint.PrettyPrinter()
 
-    def to_string(self, obj):
+    def to_string(self, obj) -> str:
         if isinstance(obj, _SubmarineObject):
             return "<%s: %s>" % (get_classname(obj), self._entity_to_string(obj))
         return self.printer.pformat(obj)
 
-    def _entity_to_string(self, entity):
+    def _entity_to_string(self, entity) -> str:
         return ", ".join(["%s=%s" % (key, self.to_string(value)) for key, value in entity])

submarine-sdk/pysubmarine/submarine/entities/experiment.py

@@ -13,6 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Type
+
+from sqlalchemy.sql.schema import Column
+
 from submarine.entities._submarine_object import _SubmarineObject
 
 
@@ -21,7 +25,15 @@ class Experiment(_SubmarineObject):
     Experiment object.
     """
 
-    def __init__(self, id, experiment_spec, create_by, create_time, update_by, update_time):
+    def __init__(
+        self,
+        id: Type[Column],
+        experiment_spec: Type[Column],
+        create_by: Type[Column],
+        create_time: Type[Column],
+        update_by: Type[Column],
+        update_time: Type[Column],
+    ):
         self._id = id
         self._experiment_spec = experiment_spec
         self._create_by = create_by
@@ -30,31 +42,31 @@ def __init__(self, id, experiment_spec, create_by, create_time, update_by, updat
         self._update_time = update_time
 
     @property
-    def id(self):
+    def id(self) -> Type[Column]:
         """String ID of the experiment."""
         return self._id
 
     @property
-    def experiment_spec(self):
+    def experiment_spec(self) -> Type[Column]:
         """String of the experiment spec."""
         return self._experiment_spec
 
     @property
-    def create_by(self):
+    def create_by(self) -> Type[Column]:
         """String name of created user id."""
         return self.create_by
 
     @property
-    def create_time(self):
+    def create_time(self) -> Type[Column]:
         """Datetime of create time."""
         return self._create_time
 
     @property
-    def update_by(self):
+    def update_by(self) -> Type[Column]:
         """String name of updated user id"."""
         return self._update_by
 
     @property
-    def update_time(self):
+    def update_time(self) -> Type[Column]:
         """Datetime of update time."""
         return self._update_time

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

@@ -26,7 +26,7 @@
 _CANONICAL_MAPPING = {stage.lower(): stage for stage in ALL_STAGES}
 
 
-def get_canonical_stage(stage):
+def get_canonical_stage(stage: str):
     key = stage.lower()
     if key not in _CANONICAL_MAPPING:
         raise SubmarineException(f"Invalid Model Version stage {stage}.")

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

@@ -13,6 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Type
+
+from sqlalchemy.sql.schema import Column
+
 from submarine.entities._submarine_object import _SubmarineObject
 
 
@@ -23,14 +27,14 @@ class ModelVersion(_SubmarineObject):
 
     def __init__(
         self,
-        name,
-        version,
-        source,
-        user_id,
-        experiment_id,
-        current_stage,
-        creation_time,
-        last_updated_time,
+        name: Type[Column],
+        version: Type[Column],
+        source: Type[Column],
+        user_id: Type[Column],
+        experiment_id: Type[Column],
+        current_stage: Type[Column],
+        creation_time: Type[Column],
+        last_updated_time: Type[Column],
         dataset=None,
         description=None,
         tags=None,

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

@@ -13,6 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Type
+
+from sqlalchemy.sql.schema import Column
+
 from submarine.entities._submarine_object import _SubmarineObject
 
 
@@ -21,7 +25,7 @@ class ModelVersionTag(_SubmarineObject):
     Tag object associated with a model version.
     """
 
-    def __init__(self, tag):
+    def __init__(self, tag: Type[Column]):
         self._tag = tag
 
     @property

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

@@ -13,6 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Type
+
+from sqlalchemy.sql.schema import Column
+
 from submarine.entities._submarine_object import _SubmarineObject
 
 
@@ -21,7 +25,14 @@ class RegisteredModel(_SubmarineObject):
     Registered model object.
     """
 
-    def __init__(self, name, creation_time, last_updated_time, description=None, tags=None):
+    def __init__(
+        self,
+        name: Type[Column],
+        creation_time: Type[Column],
+        last_updated_time: Type[Column],
+        description=None,
+        tags=None,
+    ):
         self._name = name
         self._creation_time = creation_time
         self._last_updated_time = last_updated_time

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

@@ -13,6 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Type
+
+from sqlalchemy.sql.schema import Column
+
 from submarine.entities._submarine_object import _SubmarineObject
 
 
@@ -21,7 +25,7 @@ class RegisteredModelTag(_SubmarineObject):
     Tag object associated with a registered model.
     """
 
-    def __init__(self, tag):
+    def __init__(self, tag: Type[Column]):
         self._tag = tag
 
     @property

submarine-sdk/pysubmarine/submarine/experiment/api/experiment_client.py

@@ -38,7 +38,7 @@ def generate_host():
 
 
 class ExperimentClient:
-    def __init__(self, host=generate_host()):
+    def __init__(self, host: str = generate_host()):
         """
         Submarine experiment client constructor
         :param host: An HTTP URI like http://submarine-server:8080.
@@ -59,7 +59,7 @@ def create_experiment(self, experiment_spec):
         response = self.experiment_api.create_experiment(experiment_spec=experiment_spec)
         return response.result
 
-    def wait_for_finish(self, id, polling_interval=10):
+    def wait_for_finish(self, id, polling_interval: float = 10):
         """
         Waits until experiment is finished or failed
         :param id: submarine experiment id
@@ -75,7 +75,7 @@ def wait_for_finish(self, id, polling_interval=10):
             index = self._log_pod(id, index)
             time.sleep(polling_interval)
 
-    def _log_pod(self, id, index):
+    def _log_pod(self, id, index: int):
         response = self.experiment_api.get_log(id)
         log_contents = response.result["logContent"]
         if len(log_contents) == 0: