From 2588cc2eece903d586595f16e7b5f7e2078d5c01 Mon Sep 17 00:00:00 2001
From: polyaxon-ci <ci@polyaxon.com>
Date: Thu, 26 Oct 2023 12:54:23 +0200
Subject: [PATCH] Auto-persist configs for sandboxes

---
 haupt/haupt/cli/runners/sandbox.py     |  2 +-
 haupt/haupt/common/settings/core.py    | 14 +++++++++-----
 haupt/haupt/schemas/platform_config.py |  2 +-
 haupt/haupt/schemas/sandbox_config.py  | 10 ++++++++++
 haupt/haupt/settings.py                |  7 +++++--
 5 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/haupt/haupt/cli/runners/sandbox.py b/haupt/haupt/cli/runners/sandbox.py
index a875a347d1..36dec55b3c 100644
--- a/haupt/haupt/cli/runners/sandbox.py
+++ b/haupt/haupt/cli/runners/sandbox.py
@@ -19,7 +19,7 @@ def start(
     uds: Optional[str] = None,
 ):
     os.environ[ENV_KEYS_SERVICE] = PolyaxonServices.API
-    settings.set_sandbox_config(path=ctx_paths.CONTEXT_ARTIFACTS_ROOT)
+    settings.set_sandbox_config(path=ctx_paths.CONTEXT_ARTIFACTS_ROOT, persist=True)
     host = host or settings.SANDBOX_CONFIG.host
     port = port or settings.SANDBOX_CONFIG.port
     start_app(
diff --git a/haupt/haupt/common/settings/core.py b/haupt/haupt/common/settings/core.py
index 71e55f8bd3..8ae0114063 100644
--- a/haupt/haupt/common/settings/core.py
+++ b/haupt/haupt/common/settings/core.py
@@ -40,11 +40,15 @@ def set_core(context, config: PlatformConfig, use_db: bool = True):
         context["DEFAULT_DB_ENGINE"] = db_engine
         db_name = config.db_name
         if not db_name:
-            db_name = (
-                f"{settings.SANDBOX_CONFIG.path}/plxdb"
-                if config.is_sqlite_db_engine
-                else "polyaxon"
-            )
+            if config.is_sqlite_db_engine:
+                db_path = (
+                    settings.SANDBOX_CONFIG.path
+                    if settings.SANDBOX_CONFIG and settings.SANDBOX_CONFIG.path
+                    else "/tmp"
+                )
+                db_name = f"{db_path}/plx.db"
+            else:
+                db_name = "polyaxon"
         db_definition = {
             "ENGINE": db_engine,
             "NAME": db_name,
diff --git a/haupt/haupt/schemas/platform_config.py b/haupt/haupt/schemas/platform_config.py
index 7ffd178859..23b1adf8fc 100644
--- a/haupt/haupt/schemas/platform_config.py
+++ b/haupt/haupt/schemas/platform_config.py
@@ -53,7 +53,7 @@ class PlatformConfig(BaseSchemaModel):
     logs_root: Optional[str] = Field(
         alias=ENV_KEYS_LOGS_ROOT, default=DEFAULT_LOGS_ROOT
     )
-    log_level: Optional[str] = Field(alias=ENV_KEYS_LOG_LEVEL, default="WARNING")
+    log_level: Optional[str] = Field(alias=ENV_KEYS_LOG_LEVEL, default="INFO")
     timezone: Optional[str] = Field(alias=ENV_KEYS_TIME_ZONE, default="UTC")
     scheduler_enabled: Optional[bool] = Field(
         alias="POLYAXON_SCHEDULER_ENABLED", default=False
diff --git a/haupt/haupt/schemas/sandbox_config.py b/haupt/haupt/schemas/sandbox_config.py
index 54b9f2fef1..ce1840a369 100644
--- a/haupt/haupt/schemas/sandbox_config.py
+++ b/haupt/haupt/schemas/sandbox_config.py
@@ -1,7 +1,10 @@
+import logging
+
 from typing import Optional
 
 from clipped.compact.pydantic import Field, StrictInt, StrictStr, root_validator
 from clipped.utils.http import clean_host
+from clipped.utils.paths import check_or_create_path
 
 from polyaxon._env_vars.keys import (
     ENV_KEYS_K8S_NAMESPACE,
@@ -16,6 +19,8 @@
 )
 from polyaxon._schemas.agent import AgentConfig
 
+_logger = logging.getLogger("sandbox.config")
+
 
 class SandboxConfig(AgentConfig):
     _IDENTIFIER = "sandbox"
@@ -68,3 +73,8 @@ def mount_sandbox(self, path: Optional[str] = None):
         from polyaxon._contexts.paths import mount_sandbox
 
         self.path = mount_sandbox(path=self.path or path)
+        try:
+            check_or_create_path(self.path, is_dir=True)
+        except Exception as e:
+            _logger.error("Could not create sandbox path `%s`.", self.path)
+            raise e
diff --git a/haupt/haupt/settings.py b/haupt/haupt/settings.py
index a009e531fb..50051d16d3 100644
--- a/haupt/haupt/settings.py
+++ b/haupt/haupt/settings.py
@@ -19,7 +19,7 @@ def set_proxies_config():
     PROXIES_CONFIG = ProxiesManager.get_config_from_env()
 
 
-def set_sandbox_config(path: Optional[str] = None):
+def set_sandbox_config(path: Optional[str] = None, persist: bool = False):
     from haupt.managers.sandbox import SandboxConfigManager
     from polyaxon.settings import HOME_CONFIG, set_agent_config
 
@@ -32,7 +32,10 @@ def set_sandbox_config(path: Optional[str] = None):
         SANDBOX_CONFIG = SandboxConfigManager.get_config_or_default()
         SANDBOX_CONFIG.mount_sandbox(path=path)
         SANDBOX_CONFIG.set_default_artifacts_store()
-        set_agent_config(SANDBOX_CONFIG)
+        if persist:
+            SandboxConfigManager.set_config(SANDBOX_CONFIG)
+            set_agent_config(SANDBOX_CONFIG, persist=True)
+
     except (TypeError, ValidationError):
         SandboxConfigManager.purge()
         Printer.warning("Your sandbox configuration was purged!")