Support pg8000 to support Google cloud sql (#76)

* Support pg8000 to Google support

Signed-off-by: kerthcet <[email protected]>

* fix lint

Signed-off-by: kerthcet <[email protected]>

---------

Signed-off-by: kerthcet <[email protected]>

Author: kerthcet

.env.example

@@ -1,4 +1,6 @@
 # We use PG by default as the metadata database.
 METADATA_DB_URL=postgresql+psycopg2://alphatrion:alphatr1on@localhost:5432/alphatrion
-ARTIFACT_REGISTRY_URL=http://localhost:5000/
 LOG_LEVEL=INFO
+# Default: true, ARTIFACT_REGISTRY_URL should not be empty.
+ENABLE_ARTIFACT_STORAGE=true
+ARTIFACT_REGISTRY_URL=http://localhost:5000/

alphatrion/consts.py

@@ -1,2 +1,3 @@
 METADATA_DB_URL = "METADATA_DB_URL"
 ARTIFACT_REGISTRY_URL = "ARTIFACT_REGISTRY_URL"
+ENABLE_ARTIFACT_STORAGE = "ENABLE_ARTIFACT_STORAGE"

alphatrion/log/log.py

@@ -27,6 +27,12 @@ async def log_artifact(
     if runtime is None:
         raise RuntimeError("Runtime is not initialized. Please call init() first.")
 
+    if not runtime.artifact_storage_enabled():
+        raise RuntimeError(
+            "Artifact storage is not enabled in the runtime."
+            "Set ENABLE_ARTIFACT_STORAGE=true in the environment variables."
+        )
+
     # We use experiment ID as the repo name rather than the experiment name,
     # because experiment name is not unique
     exp = runtime.current_exp

alphatrion/runtime/runtime.py

@@ -51,9 +51,14 @@ def __init__(
         self._metadb = SQLStore(
             os.getenv(consts.METADATA_DB_URL), init_tables=init_tables
         )
-        self._artifact = Artifact(
-            project_id=self._project_id, insecure=artifact_insecure
-        )
+
+        if self.artifact_storage_enabled():
+            self._artifact = Artifact(
+                project_id=self._project_id, insecure=artifact_insecure
+            )
+
+    def artifact_storage_enabled(self) -> bool:
+        return os.getenv(consts.ENABLE_ARTIFACT_STORAGE, "true").lower() == "true"
 
     # current_exp is the current running experiment.
     @property

pyproject.toml

@@ -15,6 +15,7 @@ dependencies = [
     "traceloop-sdk (>=0.47.5,<1.0.0)",
     "openai (>=2.6.1,<3.0.0)",
     "alembic (>=1.17.1,<2.0.0)",
+    "pg8000>=1.31.5",
 ]
 
 [dependency-groups]