Aurora job with Nix packages example

elasticsearch.aurora.py

# Most of a working example elasticsearch instance

include('common/nix.aurora')
include('common/util.aurora')

proc_elasticsearch = Process(
    name = 'elasticsearch',
    daemon = True,
    cmdline = unindent(
        r"""
        nix-channel --add https://my.hydra.daemon.example.com/jobset/production/master/channel/latest mycompany
        nix-channel --update
        nix-env -iA mycompany.elasticsearch
        
        mkdir -p config
        cat > config/logging.yml <<EOF
        log4j stuff blah blah blah
        EOF
        
        export ES_HOME=$PWD ES_MIN_MEM=512m ES_MAX_MEM=4096m
        exec elasticsearch \
            --node.name "$HOSTNAME" \
            --cluster.name "{{es_cluster}}" \
            --bootstrap.mlockall true \
            --http.ports "{{thermos.ports[http]}}" \
            --transport.tcp.port "{{thermos.ports[transport]}}"
        """))

task_elasticsearch = Task(
    name = "elasticsearch",
    resources = Resources(ram=4*GB, cpu=1.0, disk=25*GB),
    processes = [proc_nixinit, proc_elasticsearch, proc_nixcleanup],
    constraints = order(proc_nixinit, proc_elasticsearch))

job_es = Job(
    name = 'elasticsearch',
    service = True,
    announce = Announcer(primary_port='http'))

jobs = [ job_es(cluster=..., role=..., environment=...) ]

nix.aurora.py

# Nix-related helpers
include('util.aurora')

# This could be simplified in various ways:
# - You might not want to include root's channels.
# - In theory the environment variables related to SSL certs can be managed better elsewhere
# - If you can use nix without relying on nix-env to manage environments, such as
#   using nix-build, nix-shell, and/or nix-exec in creative ways, you could
#   skip all the parts related to profiles, which means you could also drop the
#   cleanup process.

class Nix(Struct):
    env_path = Default(String, ':'.join(['$HOME/.nix-profile/bin',
                                         '$HOME/.nix-profile/sbin',
                                         '/nix/var/nix/profiles/default/bin',
                                         '/nix/var/nix/profiles/default/sbin',
                                         '$PATH']))

    profile_path = Default(String,
                           "/nix/var/nix/profiles/${USER}/{{thermos.task_id}}")

    preamble = Default(String, unindent(
        r"""
        if [ ! -e "$HOME/.nix-defexpr" ]; then
          mkdir -p "$HOME/.nix-defexpr"
          ln -s /nix/var/nix/profiles/per-user/root/channels "$HOME/.nix-defexpr/channels_root"
        fi
        export NIX_PATH=/nix/var/nix/profiles/per-user/root/channels
        export NIX_REMOTE=daemon
        export PATH="{{env_path}}"
        export OPENSSL_X509_CERT_FILE="/etc/ssl/certs/ca-certificates.crt"
        export GIT_SSL_CAINFO="/etc/ssl/certs/ca-certificates.crt"
        export CURL_CA_BUNDLE="/etc/ssl/certs/ca-certificates.crt"
        if [ ! -e "$HOME/.nix-profile" ]; then
          mkdir -p $(dirname "{{profile_path}}")
          nix-env --switch-profile "{{profile_path}}"
        fi
        """))

# Here is a more minimal preamble to use if you:
# - forego the use of nix-env
# - skip system-wide channels, and add them per-job or get nix expressions
#   by some other means
# - manage SSL CA bundle stuff elsewhere (I'm sure it can be done cleanly with
#   nix but I don't remember the details)
#
# If you were to install the nix commands to /usr/bin/nix-* or /usr/local/bin/*
# you could probably even skip the PATH manipulation.
preamble_minimal = unindent(
    r"""
    export NIX_REMOTE=daemon
    export PATH="/nix/var/nix/profiles/default/bin:/nix/var/nix/profiles/default/sbin:$PATH"
    """))

NixProcess = Process().bind(Nix())

# Run this process first in your task:
proc_nixinit = NixProcess(
    name = 'nix_thermos_profile',
    cmdline = unindent(
        r"""
        if [[ -e .thermos_profile ]]; then
          echo "Appending to .thermos_profile"
        else
          echo "Creating .thermos_profile"
        fi
        cat >> .thermos_profile <<'EOF'
        {{preamble}}
        EOF
        """))

# Run this last:
proc_nixcleanup = NixProcess(
    name = 'nix_cleanup',
    final = True,
    cmdline = unindent(
        r"""
        # Remove this job's nix profile so it can be garbage collected later
        rm -rfv -- "{{profile_path}}"
        """))

util.aurora.py

# Handy helpers for aurora files

import json
import textwrap
#import yaml

# De-indent a string and strip leading/trailing whitespace.
unindent = lambda x: textwrap.dedent(x).strip()

# Generate varname='varvalue' shell environment assignments from a dictionary
genShellEnv = lambda **x: " \\\n  ".join("%s='%s'" % (k, v) for k, v in x.items())

to_json = lambda x: json.dumps(x, indent=2)

# the aurora2 client doesn't seem to include libyaml.
#to_yaml = lambda x: yaml.dumps(x, indent=2)

# Like SimpleTask, but for things that don't need 1gb ram and a whole cpu.
def TinyTask(name, command):
    return Task(
        resources=Resources(ram = 512*MB, cpu = 0.1, disk = 100*MB),
        processes=[Process(name=name, cmdline=command)])