No module named psycopg2 when using postgresql

[nitish-predera]

args: ["server", "--host", "0.0.0.0", "-p", "80", "--backend-store-uri", "postgresql://mlflow:mlflow@postgres-postgresql:5432/mlflow", "--default-artifact-root", "s3://mlflow"]

Adding the above parameters to use postgresql as backend store gave me the following error.

2020/05/08 04:51:31 ERROR mlflow.cli: Error initializing backend store
2020/05/08 04:51:31 ERROR mlflow.cli: No module named 'psycopg2'
Traceback (most recent call last):
  File "/usr/local/lib/python3.8/site-packages/mlflow/cli.py", line 286, in server
    initialize_backend_stores(backend_store_uri, default_artifact_root)
  File "/usr/local/lib/python3.8/site-packages/mlflow/server/handlers.py", line 101, in initialize_backend_stores
    _get_tracking_store(backend_store_uri, default_artifact_root)
  File "/usr/local/lib/python3.8/site-packages/mlflow/server/handlers.py", line 87, in _get_tracking_store
    _tracking_store = _tracking_store_registry.get_store(store_uri, artifact_root)
  File "/usr/local/lib/python3.8/site-packages/mlflow/tracking/_tracking_service/registry.py", line 37, in get_store
    return builder(store_uri=store_uri, artifact_uri=artifact_uri)
  File "/usr/local/lib/python3.8/site-packages/mlflow/server/handlers.py", line 60, in _get_sqlalchemy_store
    return SqlAlchemyStore(store_uri, artifact_uri)
  File "/usr/local/lib/python3.8/site-packages/mlflow/store/tracking/sqlalchemy_store.py", line 84, in __init__
    self.engine = mlflow.store.db.utils.create_sqlalchemy_engine(db_uri)
  File "/usr/local/lib/python3.8/site-packages/mlflow/store/db/utils.py", line 187, in create_sqlalchemy_engine
    return sqlalchemy.create_engine(db_uri, pool_pre_ping=True,
  File "/usr/local/lib/python3.8/site-packages/sqlalchemy/engine/__init__.py", line 479, in create_engine
    return strategy.create(*args, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/sqlalchemy/engine/strategies.py", line 87, in create
    dbapi = dialect_cls.dbapi(**dbapi_args)
  File "/usr/local/lib/python3.8/site-packages/sqlalchemy/dialects/postgresql/psycopg2.py", line 737, in dbapi
    import psycopg2
ModuleNotFoundError: No module named 'psycopg2'

[jadia]

Adding

RUN apt update && \ 
         apt install -y gcc libpq-dev && \
         pip install psycopg2==2.8.5

to the Dockerfile resolves the error.
Shall I put a PR for this?

[tedostrem]

yes please do, thank you !

[abedinia]

hello, would you please merge this PR? it will solve the problem.
Pull-Request

[Sairam90]

Could some one explain how to fix this?