Postgres with db

cmflib/cmf.py

@@ -25,7 +25,6 @@
 
 # This import is needed for jupyterlab environment
 from ml_metadata.proto import metadata_store_pb2 as mlpb
-from ml_metadata.metadata_store import metadata_store
 from cmflib.dvc_wrapper import (
     dvc_get_url,
     dvc_get_hash,
@@ -40,6 +39,8 @@
     git_commit,
 )
 from cmflib import graph_wrapper
+from cmflib.store.sqllite_store import SqlliteStore
+from cmflib.store.postgres import PostgresStore 
 from cmflib.metadata_helper import (
     get_or_create_parent_context,
     get_or_create_run_context,
@@ -54,7 +55,7 @@
     link_execution_to_input_artifact,
 )
 from cmflib.utils.cmf_config import CmfConfig
-from cmflib.utils.helper_functions import get_python_env, change_dir
+from cmflib.utils.helper_functions import change_dir
 from cmflib.cmf_commands_wrapper import (
     _metadata_push,
     _metadata_pull,
@@ -103,16 +104,9 @@ class Cmf:
     """
 
     # pylint: disable=too-many-instance-attributes
-    # Reading CONFIG_FILE variable
-    cmf_config = os.environ.get("CONFIG_FILE", ".cmfconfig")
     ARTIFACTS_PATH = "cmf_artifacts"
     DATASLICE_PATH = "dataslice"
     METRICS_PATH = "metrics"
-    if os.path.exists(cmf_config):
-        attr_dict = CmfConfig.read_config(cmf_config)
-        __neo4j_uri = attr_dict.get("neo4j-uri", "")
-        __neo4j_password = attr_dict.get("neo4j-password", "")
-        __neo4j_user = attr_dict.get("neo4j-user", "")
 
     def __init__(
         self,
@@ -128,17 +122,30 @@ def __init__(
 					else  os.getcwd()
 
         logging_dir = change_dir(self.cmf_init_path)
+        temp_store = ""
         if is_server is False:
             Cmf.__prechecks()
+            temp_store = SqlliteStore({"filename":filepath})
+        else:
+            IP = os.getenv('MYIP')
+            POSTGRES_DB = os.getenv('POSTGRES_DB')
+            POSTGRES_USER = os.getenv('POSTGRES_USER')
+            POSTGRES_PASSWORD = os.getenv('POSTGRES_PASSWORD')
+            #print(f"The value of POSTGRES_DB is {POSTGRES_DB}")
+            #print(f"The value of POSTGRES_USER: {POSTGRES_USER}")
+            #print(f"The value of POSTGRES_PASSSWORD: {POSTGRES_PASSWORD}")
+            #print(f"The value of POSTGRES_HOST: {IP}")
+            config_dict = {"host":IP, "port":"5432", "user": POSTGRES_USER, "password": POSTGRES_PASSWORD, "dbname": POSTGRES_DB}
+            temp_store = PostgresStore(config_dict)
+        #print("temp_store type", type(temp_store))
         if custom_properties is None:
             custom_properties = {}
         if not pipeline_name:
             # assign folder name as pipeline name 
             cur_folder = os.path.basename(os.getcwd())
             pipeline_name = cur_folder
-        config = mlpb.ConnectionConfig()
-        config.sqlite.filename_uri = filepath
-        self.store = metadata_store.MetadataStore(config)
+        self.store = temp_store.connect()
+        #print("self.store = ", self.store)
         self.filepath = filepath
         self.child_context = None
         self.execution = None
@@ -419,7 +426,7 @@ def create_execution(
         git_repo = git_get_repo()
         git_start_commit = git_get_commit()
         cmd = str(sys.argv) if cmd is None else cmd
-        python_env=get_python_env()
+
         self.execution = create_new_execution_in_existing_run_context(
             store=self.store,
             # Type field when re-using executions
@@ -433,15 +440,14 @@ def create_execution(
             pipeline_type=self.parent_context.name,
             git_repo=git_repo,
             git_start_commit=git_start_commit,
-            python_env=python_env,
             custom_properties=custom_props,
             create_new_execution=create_new_execution,
         )
         uuids = self.execution.properties["Execution_uuid"].string_value
         if uuids:
             self.execution.properties["Execution_uuid"].string_value = uuids+","+str(uuid.uuid1())
         else:
-            self.execution.properties["Execution_uuid"].string_value = str(uuid.uuid1())            
+            self.execution.properties["Execution_uuid"].string_value = str(uuid.uuid1())          
         self.store.put_executions([self.execution])
         self.execution_name = str(self.execution.id) + "," + execution_type
         self.execution_command = cmd
@@ -451,7 +457,7 @@ def create_execution(
         self.execution_label_props["Execution_Name"] = (
             execution_type + ":" + str(self.execution.id)
         )
-        
+
         self.execution_label_props["execution_command"] = cmd
         if self.graph:
             self.driver.create_execution_node(
@@ -602,7 +608,6 @@ def merge_created_execution(
         # print(custom_props)
         git_repo = properties.get("Git_Repo", "")
         git_start_commit = properties.get("Git_Start_Commit", "")
-        python_env = properties.get("Python_Env", "")
         #name = properties.get("Name", "")
         create_new_execution = True
         execution_name = execution_type
@@ -623,7 +628,6 @@ def merge_created_execution(
             pipeline_type=self.parent_context.name,
             git_repo=git_repo,
             git_start_commit=git_start_commit,
-            python_env=python_env,
             custom_properties=custom_props,
             create_new_execution=create_new_execution
         )
@@ -658,13 +662,15 @@ def merge_created_execution(
                 self.execution.id,
                 custom_props,
             )
+
         return self.execution
 
     def log_dvc_lock(self, file_path: str):
         """Used to update the dvc lock file created with dvc run command."""
         print("Entered dvc lock file commit")
         return commit_dvc_lock_file(file_path, self.execution.id)
 
+
     def log_dataset(
         self,
         url: str,
@@ -1996,6 +2002,7 @@ def commit_existing(self, uri: str, props: t.Optional[t.Dict] = None, custom_pro
 #                print(last)
 #                os.symlink(str(index), slicedir + "/ " + last)
 
+
 def metadata_push(pipeline_name: str, filepath = "./mlmd", tensorboard_path: str = "", execution_id: str = ""):
     """ Pushes MLMD file to CMF-server.
     Example:

cmflib/cmf_merger.py

@@ -165,18 +165,18 @@ def create_original_time_since_epoch(mlmd_data):
         "Pipeline"
     ][0]["create_time_since_epoch"]
     for i in mlmd_data["Pipeline"][0]["stages"]:
-        i["custom_properties"]["original_create_time_since_epoch"] = i[
+        i["custom_properties"]["original_create_time_since_epoch"] = str(i[
             "create_time_since_epoch"
-        ]
+        ])
         original_stages.append(
             i["custom_properties"]["original_create_time_since_epoch"]
         )
         stages.append(i["create_time_since_epoch"])
         # print(i['custom_properties']['original_create_time_since_epoch'])
         for j in i["executions"]:
-            j["custom_properties"]["original_create_time_since_epoch"] = j[
+            j["custom_properties"]["original_create_time_since_epoch"] = str(j[
                 "create_time_since_epoch"
-            ]
+            ])
             original_execution.append(
                 j["custom_properties"]["original_create_time_since_epoch"]
             )
@@ -185,7 +185,7 @@ def create_original_time_since_epoch(mlmd_data):
             for k in j["events"]:
                 k["artifact"]["custom_properties"][
                     "original_create_time_since_epoch"
-                ] = k["artifact"]["create_time_since_epoch"]
+                ] = str(k["artifact"]["create_time_since_epoch"])
                 original_artifact.append(
                     k["artifact"]["custom_properties"][
                         "original_create_time_since_epoch"
@@ -197,4 +197,3 @@ def create_original_time_since_epoch(mlmd_data):
     return mlmd_data
 
 
-

cmflib/cmfquery.py

@@ -14,13 +14,16 @@
 # limitations under the License.
 ###
 import abc
+import os
 import json
 import logging
 import typing as t
 from enum import Enum
 from google.protobuf.json_format import MessageToDict
 import pandas as pd
-from ml_metadata.metadata_store import metadata_store
+#from ml_metadata.metadata_store import metadata_store
+from cmflib.store.sqllite_store import SqlliteStore
+from cmflib.store.postgres import PostgresStore
 from ml_metadata.proto import metadata_store_pb2 as mlpb
 from cmflib.mlmd_objects import CONTEXT_LIST
 
@@ -113,10 +116,23 @@ class CmfQuery(object):
         filepath: Path to the MLMD database file.
     """
 
-    def __init__(self, filepath: str = "mlmd") -> None:
-        config = mlpb.ConnectionConfig()
-        config.sqlite.filename_uri = filepath
-        self.store = metadata_store.MetadataStore(config)
+    def __init__(self, filepath: str = "mlmd", is_server=False) -> None:
+        if is_server:
+            IP = os.getenv('MYIP')
+            POSTGRES_DB = os.getenv('POSTGRES_DB')
+            POSTGRES_USER = os.getenv('POSTGRES_USER')
+            POSTGRES_PASSWORD = os.getenv('POSTGRES_PASSWORD')
+            #print(f"The value of POSTGRES_DB is {POSTGRES_DB}")
+            #print(f"The value of POSTGRES_USER: {POSTGRES_USER}")
+            #print(f"The value of POSTGRES_PASSSWORD: {POSTGRES_PASSWORD}")
+            #print(f"The value of POSTGRES_HOST: {IP}")
+            config_dict = {"host":IP, "port":"5432", "user": POSTGRES_USER, "password": POSTGRES_PASSWORD, "dbname": POSTGRES_DB}
+            temp_store = PostgresStore(config_dict)
+        else:
+            temp_store = SqlliteStore({"filename":filepath})
+        #print("temp_store type", type(temp_store))
+        self.store = temp_store.connect()
+        #print("self.store = ", self.store)
 
     @staticmethod
     def _copy(

cmflib/store/__init__.py

@@ -0,0 +1 @@
+

cmflib/store/cmfstore.py

@@ -0,0 +1,13 @@
+from abc import ABC, abstractmethod
+from ml_metadata.metadata_store import metadata_store
+
+class CmfStore(ABC):
+
+    def __init__(self, config):
+        self.config = config
+        super().__init__()
+
+    @abstractmethod
+    def connect(self)-> metadata_store:
+        cmf_store = metadata_store.MetadataStore(self.config)
+        return cmf_store

cmflib/store/postgres.py

@@ -0,0 +1,20 @@
+#from cmfstore import CmfStore
+from cmflib.store.cmfstore import CmfStore
+from ml_metadata.proto import metadata_store_pb2 as mlpb
+from ml_metadata.metadata_store import metadata_store
+
+class PostgresStore(CmfStore):
+
+    def __init__(self, config):
+        self.connection_config = mlpb.ConnectionConfig()
+        self.connection_config.postgresql.host = config["host"]
+        self.connection_config.postgresql.port = config["port"]
+        self.connection_config.postgresql.user = config["user"]
+        self.connection_config.postgresql.password = config["password"]
+        self.connection_config.postgresql.dbname = config["dbname"]
+        super().__init__(self.connection_config)
+
+    def connect(self)-> metadata_store:
+        cmf_store = super().connect()
+        return cmf_store
+

cmflib/store/sqllite_store.py

@@ -0,0 +1,17 @@
+from cmflib.store.cmfstore import CmfStore
+from ml_metadata.metadata_store import metadata_store
+from ml_metadata.proto import metadata_store_pb2 as mlpb
+
+class SqlliteStore(CmfStore):
+
+    def __init__(self, config):
+
+        self.connection_config = mlpb.ConnectionConfig()
+        self.connection_config.sqlite.filename_uri = config["filename"]
+        self.connection_config.sqlite.connection_mode  = 3
+        super().__init__(self.connection_config)
+
+    def connect(self)-> metadata_store:
+        cmf_store = super().connect()
+        return cmf_store
+

cmflib/utils/helper_functions.py

@@ -131,4 +131,4 @@ def generate_osdf_token(key_id, key_path, key_issuer) -> str:
     except Exception as err:
         print(f"Unexpected {err}, {type(err)}")
 
-    return dynamic_pass
+    return dynamic_pass

docker-compose-server.yml

@@ -15,6 +15,23 @@
 ###
 
 services:
+  postgres:
+    image: postgres:13
+    container_name: postgres
+    environment:
+      POSTGRES_USER: ${POSTGRES_USER}
+      POSTGRES_PASSWORD: ${POSTGRES_PASSWORD}
+      POSTGRES_DB: ${POSTGRES_DB}
+    ports:
+      - "5432:5432"
+    volumes:
+      - /home/xxxx/cmf-server/data/postgres_data:/var/lib/postgresql/data
+    healthcheck:
+      test: ["CMD-SHELL", "pg_isready -d $${POSTGRES_DB} -U $${POSTGRES_USER}"]
+      interval: 10s
+      timeout: 5s
+      retries: 5
+
   tensorboard:
     image: tensorflow/tensorflow
     command: tensorboard --logdir /logs --host 0.0.0.0
@@ -24,6 +41,8 @@ services:
     # directory path should be updated as per user's environment
       - /home/xxxx/cmf-server/data/tensorboard-logs:/logs
     container_name: tensorboard
+
+
   server:
     image: server:latest
     # both the directory paths should be updated as per user's environment
@@ -39,6 +58,9 @@ services:
     environment:
       MYIP: ${IP:-127.0.0.1}
       HOSTNAME: ${hostname:-localhost}
+      POSTGRES_DB: ${POSTGRES_DB}
+      POSTGRES_USER: ${POSTGRES_USER}
+      POSTGRES_PASSWORD: ${POSTGRES_PASSWORD}
     ports:
       - "8080:80"
     expose:
@@ -47,6 +69,9 @@ services:
       test: if [ -z $IP ]; then curl -f http://$hostname:8080; else curl -f http://$IP:8080; fi
       interval: 15s
       retries: 32
+    depends_on:
+      postgres:
+        condition: service_healthy
 
   ui:
     image: ui:latest
@@ -65,3 +90,4 @@ services:
     depends_on:
       server:
         condition: service_healthy
+

server/Dockerfile

@@ -39,9 +39,6 @@ RUN apt-get update -y && apt-get install -y curl
 # Copy the requirements.txt file from the server to the cmf-server.
 COPY ./server/requirements.txt /cmf-server/requirements.txt
 
-# library required for lineage
-RUN apt-get install -y graphviz graphviz-dev 
-
 # Install requirements.txt
 RUN pip install --no-cache-dir --upgrade -r /cmf-server/requirements.txt
 

server/app/dbconfig.py

@@ -0,0 +1,14 @@
+import os
+from dotenv import load_dotenv
+
+# Load .env variables
+load_dotenv()
+
+# Database configuration
+DB_CONFIG = {
+    "user": os.getenv("POSTGRES_USER"),
+    "password": os.getenv("POSTGRES_PASSWORD"),
+    "database": os.getenv("POSTGRES_DB"),
+    "host": os.getenv("POSTGRES_HOST"),
+    "port": os.getenv("POSTGRES_PORT"),
+}