Fixes and optimization for getting db_provider.

studioml · Jul 3, 2020 · a89e9e4 · a89e9e4
1 parent d6b22ec
commit a89e9e4
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/studio/encrypted_payload_builder.py b/studio/encrypted_payload_builder.py
@@ -214,7 +214,5 @@ def construct(self, experiment, config, packages):
 
         print(json.dumps(encrypted_payload, indent=4))
 
-        exit(0)
-
         return encrypted_payload
 
diff --git a/studio/model.py b/studio/model.py
@@ -18,7 +18,7 @@
 from .local_db_provider import LocalDbProvider
 from .s3_provider import S3Provider
 from .gs_provider import GSProvider
-from .model_setup import setup_model
+import model_setup
 from . import logs
 
 def get_config(config_file=None):
@@ -59,6 +59,11 @@ def replace_with_env(config):
                      .format(config_paths))
 
 def get_db_provider(config=None, blocking_auth=True):
+
+    db_provider = model_setup.get_db_provider()
+    if not db_provider is None:
+        return db_provider
+
     if not config:
         config = get_config()
     verbose = parse_verbosity(config.get('verbose'))
@@ -115,10 +120,10 @@ def get_db_provider(config=None, blocking_auth=True):
         artifact_store = db_provider.get_artifact_store()
 
     else:
-        _model_setup = None
+        model_setup._model_setup = None
         raise ValueError('Unknown type of the database ' + db_config['type'])
 
-    setup_model(db_provider, artifact_store)
+    model_setup.setup_model(db_provider, artifact_store)
     return db_provider
 
 def parse_verbosity(verbosity=None):