def setUp(self):
environ["JOB_ID"] = JOB_ID
submarine.set_db_uri(
"mysql+pymysql://submarine_test:password_test@localhost:3306/submarine_test"
)
self.db_uri = submarine.get_db_uri()
self.client = SubmarineClient(
db_uri=self.db_uri,
s3_registry_uri=MLFLOW_S3_ENDPOINT_URL,
)
from submarine.store.tracking.sqlalchemy_store import SqlAlchemyStore
self.store = SqlAlchemyStore(self.db_uri)
from submarine.store.model_registry.sqlalchemy_store import SqlAlchemyStore
self.model_registry = SqlAlchemyStore(self.db_uri)
# TODO: use submarine.tracking.fluent to support experiment create
with self.store.ManagedSessionMaker() as session:
instance = SqlExperiment(
id=JOB_ID,
experiment_spec='{"value": 1}',
create_by="test",
create_time=datetime.now(),
update_by=None,
update_time=None,
)
session.add(instance)
session.commit()
def delete_serve(model_name: str, model_version: int):
"""
Delete a serving model
:param model_name: Name of a registered model
:param model_version: Version of a registered model
"""
SubmarineClient().delete_serve(model_name, model_version)
def create_serve(model_name: str, model_version: int):
"""
Create serve of a model through Seldon Core
:param model_name: Name of a registered model
:param model_version: Version of a registered model
"""
SubmarineClient().create_serve(model_name, model_version)
def log_param(key, value, worker_index):
"""
Log a parameter under the current run, creating a run if necessary.
:param key: Parameter name (string)
:param value: Parameter value (string, but will be string-field if not)
:param worker_index
"""
job_name = get_job_name()
SubmarineClient().log_param(job_name, key, value, worker_index)
def log_metric(key, value, worker_index, step=None):
"""
Log a metric under the current run, creating a run if necessary.
:param key: Metric name (string).
:param value: Metric value (float). Note that some special values such as +/- Infinity may be
replaced by other values depending on the store. For example, sFor example, the
SQLAlchemy store replaces +/- Inf with max / min float values.
:param worker_index: Metric worker_index (string).
:param step: Metric step (int). Defaults to zero if unspecified.
"""
job_name = get_job_name()
SubmarineClient().log_metric(job_name, key, value, worker_index,
int(time.time() * 1000), step or 0)
def save_model(
model,
model_type: str,
registered_model_name: str = None,
input_dim: list = None,
output_dim: list = None,
):
"""
Save a model into the minio pod.
:param model_type: The type of the model.
:param model: Model.
:param registered_model_name: If none None, register model into the model registry with
this name. If None, the model only be saved in minio pod.
"""
SubmarineClient().save_model(model, model_type, registered_model_name,
input_dim, output_dim)
class TestTracking(unittest.TestCase):
def setUp(self):
environ["JOB_ID"] = JOB_ID
submarine.set_db_uri(
"mysql+pymysql://submarine_test:password_test@localhost:3306/submarine_test"
)
self.db_uri = submarine.get_db_uri()
self.client = SubmarineClient(
db_uri=self.db_uri,
s3_registry_uri=MLFLOW_S3_ENDPOINT_URL,
)
from submarine.store.tracking.sqlalchemy_store import SqlAlchemyStore
self.store = SqlAlchemyStore(self.db_uri)
from submarine.store.model_registry.sqlalchemy_store import SqlAlchemyStore
self.model_registry = SqlAlchemyStore(self.db_uri)
# TODO: use submarine.tracking.fluent to support experiment create
with self.store.ManagedSessionMaker() as session:
instance = SqlExperiment(
id=JOB_ID,
experiment_spec='{"value": 1}',
create_by="test",
create_time=datetime.now(),
update_by=None,
update_time=None,
)
session.add(instance)
session.commit()
def tearDown(self):
submarine.set_db_uri(None)
models.Base.metadata.drop_all(self.store.engine)
environ["MLFLOW_S3_ENDPOINT_URL"] = MLFLOW_S3_ENDPOINT_URL
environ["AWS_ACCESS_KEY_ID"] = "submarine_minio"
environ["AWS_SECRET_ACCESS_KEY"] = "submarine_minio"
Repository().delete_folder(f"experiment/{JOB_ID}")
Repository().delete_folder(f"registry/{REGISTERED_MODEL_NAME}")
def test_log_param(self):
submarine.log_param("name_1", "a")
# Validate params
with self.store.ManagedSessionMaker() as session:
params = session.query(SqlParam).options().filter(
SqlParam.id == JOB_ID).all()
assert params[0].key == "name_1"
assert params[0].value == "a"
assert params[0].id == JOB_ID
def test_log_metric(self):
submarine.log_metric("name_1", 5)
submarine.log_metric("name_1", 6)
# Validate params
with self.store.ManagedSessionMaker() as session:
metrics = session.query(SqlMetric).options().filter(
SqlMetric.id == JOB_ID).all()
assert len(metrics) == 2
assert metrics[0].key == "name_1"
assert metrics[0].value == 5
assert metrics[0].id == JOB_ID
assert metrics[1].value == 6
@pytest.mark.skipif(tensorflow.version.VERSION < "2.0",
reason="using tensorflow 2")
def test_save_model(self):
input_arr = tensorflow.random.uniform((1, 5))
model = LinearNNModel()
model(input_arr)
self.client.save_model(model, "tensorflow", "name_1",
REGISTERED_MODEL_NAME)
self.client.save_model(model, "tensorflow", "name_2",
REGISTERED_MODEL_NAME)
# Validate model_versions
model_versions = self.model_registry.list_model_versions(
REGISTERED_MODEL_NAME)
assert len(model_versions) == 2
assert model_versions[0].name == REGISTERED_MODEL_NAME
assert model_versions[0].version == 1
assert model_versions[1].name == REGISTERED_MODEL_NAME
assert model_versions[1].version == 2