Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Print table schema in DB initialization test (#5248)
* Print out table schemas in DB initialization check Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * rename Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * fix arg name Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * rename to run_logging_operations Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * remove COPY log.py Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * move assert Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * print schema after logging Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * fix Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * add sqlite Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * update step Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * move parse_args Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * separate step Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * pull Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * remove comments Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>
- Loading branch information
Showing
10 changed files
with
96 additions
and
46 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,6 +1,5 @@ | ||
** | ||
|
||
!dist/*.whl | ||
!log.py | ||
!init-mssql-db.sh | ||
!init-mssql-db.sql |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,2 @@ | ||
schemas | ||
mlflowdb |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Empty file.
Empty file.
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,67 @@ | ||
import os | ||
import argparse | ||
|
||
import sqlalchemy | ||
from sqlalchemy.schema import MetaData, CreateTable | ||
|
||
import mlflow | ||
from mlflow.tracking._tracking_service.utils import _TRACKING_URI_ENV_VAR | ||
|
||
|
||
class MockModel(mlflow.pyfunc.PythonModel): | ||
def load_context(self, context): | ||
pass | ||
|
||
def predict(self, context, model_input): | ||
pass | ||
|
||
|
||
def parse_args(): | ||
parser = argparse.ArgumentParser() | ||
parser.add_argument("--schema-output", required=True, help="Output path of DB schema") | ||
return parser.parse_args() | ||
|
||
|
||
def run_logging_operations(): | ||
with mlflow.start_run() as run: | ||
print("Tracking URI:", mlflow.get_tracking_uri()) | ||
mlflow.log_param("p", "param") | ||
mlflow.log_metric("m", 1.0) | ||
mlflow.set_tag("t", "tag") | ||
mlflow.pyfunc.log_model( | ||
artifact_path="model", | ||
python_model=MockModel(), | ||
registered_model_name="mock", | ||
) | ||
print(mlflow.get_run(run.info.run_id)) | ||
|
||
|
||
def get_db_schema(): | ||
engine = sqlalchemy.create_engine(mlflow.get_tracking_uri()) | ||
created_tables_metadata = MetaData(bind=engine) | ||
created_tables_metadata.reflect() | ||
# Write out table schema as described in | ||
# https://docs.sqlalchemy.org/en/13/faq/metadata_schema.html#how-can-i-get-the-create-table-drop-table-output-as-a-string | ||
lines = [] | ||
for ti in created_tables_metadata.sorted_tables: | ||
lines += list(map(str.rstrip, str(CreateTable(ti)).splitlines())) | ||
return "\n".join(lines) | ||
|
||
|
||
def main(): | ||
assert _TRACKING_URI_ENV_VAR in os.environ | ||
|
||
args = parse_args() | ||
run_logging_operations() | ||
schema = get_db_schema() | ||
title = "Schema" | ||
print("=" * 10, title, "=" * 10) | ||
print(schema) | ||
print("=" * (20 + 2 + len(title))) | ||
os.makedirs(os.path.dirname(args.schema_output), exist_ok=True) | ||
with open(args.schema_output, "w") as f: | ||
f.write(schema) | ||
|
||
|
||
if __name__ == "__main__": | ||
main() |