Optional ssh port forwarding (#134)

* Optional ssh port forwarding

* Update database.py

* File locking for ssh tunnel

* add configuration to ssh forwarding; move defaults to settings.py

---------

Co-authored-by: Nicholas Gao <nicholas.gao@tum.de>

Author: sigeisler

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
     - repo: https://github.com/pre-commit/pre-commit-hooks
-      rev: v4.5.0
+      rev: v4.6.0
       hooks:
           - id: check-case-conflict
           - id: check-toml
@@ -10,7 +10,7 @@ repos:
           - id: trailing-whitespace
     - repo: https://github.com/astral-sh/ruff-pre-commit
       # Ruff version.
-      rev: v0.3.0
+      rev: v0.3.5
       hooks:
           # Run the linter.
           - id: ruff

README.md

@@ -7,7 +7,7 @@ Keeping track of computational experiments can be annoying and failure to do so
 While workload scheduling systems such as [`Slurm`](https://slurm.schedmd.com/overview.html) make it easy to run many experiments in parallel on a cluster, it can be hard to keep track of which parameter configurations are running, failed, or completed.
 [`sacred`](https://github.com/IDSIA/sacred) is a great tool to collect and manage experiments and their results, especially when used with a [`MongoDB`](https://www.mongodb.com/). However, it is lacking integration with workload schedulers.
 
-**`SEML`** enables you to 
+**`SEML`** enables you to
 * very easily define hyperparameter search spaces using YAML files,
 * run these hyperparameter configurations on a compute cluster using `Slurm`,
 * and to track the experimental results using `sacred` and `MongoDB`.
@@ -31,7 +31,7 @@ conda install -c conda-forge seml
 ```
 Then configure your MongoDB via:
 ```bash
-seml configure  --mongodb # provide your MongoDB credentials
+seml configure
 ```
 For convenience, you may create your first **`SEML`** project via:
 ```bash
@@ -40,6 +40,17 @@ seml project init -t default new_project
 ```
 in an empty directoy. **`SEML`** will automatically create a python package for you.
 
+
+### SSH Port Forwarding
+If your MongoDB is only accessible via an SSH port forward, **`SEML`** allows you to directly configure this as well if you install the `ssh_forward` dependencies via:
+```bash
+pip install seml[ssh_forward]
+```
+It remains to configure the SSH settings:
+```bash
+seml configure --ssh_forward
+```
+
 ### Development
 If you want to develop `seml` please clone the repository and install it via
 ```bash
@@ -48,7 +59,7 @@ pip install -e .[dev]
 and install pre-commit hooks via
 ```bash
 pre-commit install
-``` 
+```
 
 ## Documentation
 Documentation is available in our [docs.md](docs.md) or via the CLI:

pyproject.toml

@@ -36,6 +36,7 @@ dependencies = [
 
 [project.optional-dependencies]
 dev = ["pytest", "ruff", "pre-commit"]
+ssh_forward = ["sshtunnel>=0.4.0", "filelock>=3.13.3"]
 
 [tool.ruff.format]
 quote-style = "single"

src/seml/__main__.py

@@ -356,27 +356,20 @@ def clean_db_command(ctx: typer.Context, yes: YesAnnotation = False):
 @restrict_collection(False)
 def configure_command(
     ctx: typer.Context,
-    all: Annotated[
+    ssh_forward: Annotated[
         bool,
         typer.Option(
-            '-a',
-            '--all',
-            help='Configure all SEML settings',
+            '-sf',
+            '--ssh-forward',
+            help='Configure SSH forwarding settings for MongoDB.',
             is_flag=True,
         ),
     ] = False,
-    mongodb: Annotated[
-        bool,
-        typer.Option(
-            help='Configure MongoDB settings',
-            is_flag=True,
-        ),
-    ] = True,
 ):
     """
     Configure SEML (database, argument completion, ...).
     """
-    configure(all=all, mongodb=mongodb)
+    configure(all=False, mongodb=True, setup_ssh_forward=ssh_forward)
 
 
 @app.command('start-jupyter')

src/seml/configure.py

@@ -1,9 +1,30 @@
 import logging
+import yaml
 
 from seml.settings import SETTINGS
 
 
-def mongodb_configure():
+def prompt_ssh_forward():
+    """
+    Prompt the user for SSH Forward settings. The output format corresponds
+    to the argument of sshtunnel.SSHTunnelForwarder.
+    """
+    from seml.console import prompt
+    
+    logging.info('Configuring SSH Forward settings.')
+    ssh_host = prompt('SSH host')
+    port = prompt('Port', default=22, type=int)
+    username = prompt('User name')
+    ssh_pkey = prompt('Path to SSH private key', default='~/.ssh/id_rsa')
+    return dict(
+        ssh_address_or_host=ssh_host,
+        ssh_port=port,
+        ssh_username=username,
+        ssh_pkey=ssh_pkey,
+    )
+
+
+def mongodb_configure(setup_ssh_forward: bool = False):
     from seml.console import prompt
 
     if SETTINGS.DATABASE.MONGODB_CONFIG_PATH.exists() and not prompt(
@@ -13,31 +34,36 @@ def mongodb_configure():
         return
     logging.info('Configuring MongoDB. Warning: Password will be stored in plain text.')
     host = prompt('MongoDB host')
-    port = prompt('Port', default='27017')
+    port = prompt('Port', default=27017, type=int)
     database = prompt('Database name')
     username = prompt('User name')
     password = prompt('Password', hide_input=True)
     file_path = SETTINGS.DATABASE.MONGODB_CONFIG_PATH
-    config_string = (
-        f'username: {username}\n'
-        f'password: {password}\n'
-        f'port: {port}\n'
-        f'database: {database}\n'
-        f'host: {host}'
+    config = dict(
+        host=host,
+        port=port,
+        database=database,
+        username=username,
+        password=password,
     )
+    if setup_ssh_forward:
+        config['ssh_config'] = prompt_ssh_forward()
+    config_string = yaml.dump(config)
     logging.info(
         f"Saving the following configuration to {file_path}:\n"
-        f"{config_string.replace(f'password: {password}', 'password: ********')}"
+        f"{config_string.replace(f'{password}', '********')}"
     )
     file_path.parent.mkdir(parents=True, exist_ok=True)
     with open(file_path, 'w') as f:
         f.write(config_string)
 
 
-def configure(all: bool = False, mongodb: bool = False):
+def configure(
+    all: bool = False, mongodb: bool = False, setup_ssh_forward: bool = False
+):
     configured_any = False
     if mongodb or all:
-        mongodb_configure()
+        mongodb_configure(setup_ssh_forward=setup_ssh_forward)
         configured_any = True
     if not configured_any:
         logging.info('Did not specify any configuration to configure')

src/seml/database.py

@@ -1,6 +1,10 @@
 import logging
+import random
+import time
 from typing import List
 
+import yaml
+
 from seml.errors import MongoDBError
 from seml.settings import SETTINGS
 from seml.utils import s_if
@@ -18,7 +22,69 @@ def get_collection(collection_name, mongodb_config=None, suffix=None):
     return db[collection_name]
 
 
-def get_mongo_client(db_name, host, port, username, password, **kwargs):
+def retried_and_locked_ssh_port_forward(
+    retries_max=SETTINGS.SSH_FORWARD.RETRIES_MAX,
+    retries_delay=SETTINGS.SSH_FORWARD.RETRIES_DELAY,
+    lock_file=SETTINGS.SSH_FORWARD.LOCK_FILE,
+    lock_timeout=SETTINGS.SSH_FORWARD.LOCK_TIMEOUT,
+    **ssh_config,
+):
+    try:
+        from sshtunnel import (
+            BaseSSHTunnelForwarderError,
+            SSHTunnelForwarder,
+            create_logger,
+        )
+    except ImportError:
+        logging.error(
+            'Opening ssh tunnel requires `sshtunnel` (e.g. `pip install sshtunnel`)'
+        )
+        exit(1)
+    try:
+        from filelock import FileLock, Timeout
+    except ImportError:
+        logging.error(
+            'Opening ssh tunnel requires `filelock` (e.g. `pip install filelock`)'
+        )
+        exit(1)
+
+    delay = retries_delay
+    error = None
+    for _ in range(retries_max):
+        try:
+            lock = FileLock(lock_file, timeout=lock_timeout)
+            with lock:
+                server = SSHTunnelForwarder(
+                    **ssh_config,
+                    logger=create_logger(logging.getLogger(), loglevel=logging.ERROR),
+                )
+                server.start()
+                return server
+        except Timeout as e:
+            error = e
+            logging.warn(f'Failed to aquire lock for ssh tunnel {lock_file}')
+        except BaseSSHTunnelForwarderError as e:
+            error = e
+            logging.warn(f'Retry establishing ssh tunnel in {delay} s')
+            # Jittered exponential retry
+            time.sleep(delay)
+            delay *= 2
+            delay += random.uniform(0, 1)
+
+    if error:
+        logging.error(f'Failed to establish ssh tunnel: {error}')
+        exit(1)
+
+
+def get_mongo_client(
+    db_name, host, port, username, password, ssh_config=None, **kwargs
+):
+    if ssh_config is not None:
+        server = retried_and_locked_ssh_port_forward(**ssh_config)
+
+        host = server.local_bind_host
+        port = server.local_bind_port
+
     import pymongo
 
     client = pymongo.MongoClient(
@@ -76,7 +142,8 @@ def get_mongodb_config(path=SETTINGS.DATABASE.MONGODB_CONFIG_PATH):
         - database name
         - username
         - password
-        - directConnection
+        - directConnection (Optional)
+        - ssh_config (Optional)
 
     Default path is $HOME/.config/seml/mongodb.config.
 
@@ -87,6 +154,15 @@ def get_mongodb_config(path=SETTINGS.DATABASE.MONGODB_CONFIG_PATH):
     database: <database_name>
     host: <host>
     directConnection: <bool> (Optional)
+    ssh_config: <dict> (Optional)
+      ssh_address_or_host: <the url of the jump host>
+      ssh_pkey: <the ssh host key>
+      ssh_username: <username for jump host>
+      retries_max: <number of retries to establish shh tunnel, default 6> (Optional)
+      retries_delay: <initial wait time for exponential retry, default 1> (Optional)
+      lock_file: <lockfile to avoid establishing ssh tunnel parallely, default `~/seml_ssh.lock`> (Optional)
+      lock_timeout: <timeout for aquiring lock, default 30> (Optional)
+      ** further arguments passed to `SSHTunnelForwarder` (see https://github.com/pahaz/sshtunnel)
 
     Returns
     -------
@@ -103,14 +179,8 @@ def get_mongodb_config(path=SETTINGS.DATABASE.MONGODB_CONFIG_PATH):
             f"MongoDB credentials could not be read at '{path}'.{config_str}"
         )
 
-    with open(path, 'r') as f:
-        for line in f.readlines():
-            # ignore lines that are empty
-            if len(line.strip()) > 0:
-                split = line.split(':')
-                key = split[0].strip()
-                value = split[1].strip()
-                access_dict[key] = value
+    with open(path, 'r') as conf:
+        access_dict = yaml.safe_load(conf)
 
     required_entries = ['username', 'password', 'port', 'host', 'database']
     for entry in required_entries:
@@ -129,7 +199,7 @@ def get_mongodb_config(path=SETTINGS.DATABASE.MONGODB_CONFIG_PATH):
         else False
     )
 
-    return {
+    cfg = {
         'password': db_password,
         'username': db_username,
         'host': db_host,
@@ -138,6 +208,15 @@ def get_mongodb_config(path=SETTINGS.DATABASE.MONGODB_CONFIG_PATH):
         'directConnection': db_direct,
     }
 
+    if 'ssh_config' not in access_dict:
+        return cfg
+
+    cfg['ssh_config'] = access_dict['ssh_config']
+    cfg['ssh_config']['remote_bind_address'] = (db_host, db_port)
+    cfg['directConnection'] = True
+
+    return cfg
+
 
 def build_filter_dict(filter_states, batch_id, filter_dict, sacred_id=None):
     """
@@ -255,6 +334,7 @@ def upload_file(filename, db_collection, batch_id, filetype):
 
 def delete_files(database, file_ids, progress=False):
     import gridfs
+
     from seml.console import track
 
     fs = gridfs.GridFS(database)

src/seml/settings.py

@@ -145,6 +145,12 @@
         'AUTOCOMPLETE_CACHE_ALIVE_TIME': 300,
         'SETUP_COMMAND': '',
         'END_COMMAND': '',
+        'SSH_FORWARD': {
+            'LOCK_FILE': '/tmp/seml_ssh_forward.lock',
+            'RETRIES_MAX': 6,
+            'RETRIES_DELAY': 1,
+            'LOCK_TIMEOUT': 30,
+        },
     },
 )