oarphcli

#!/usr/bin/env python
# vim: tabstop=2 shiftwidth=2 expandtab

# Copyright 2023 Maintainers of OarphPy
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


DESC = """
oarphcli - This tool serves to both automate and document the oarphpy
development workflow.

## Use & Development

$ ./oarphcli --shell
Drop into a dockerized dev environment shell, do some work. Local code mounted
at /opt/oarphpy, and outer filesystem mounted at /outer_root .


## Testing

$$ pytest oarphpy_test
$$ python3 setup.py test
In the dockerized shell, run unit tests using your local code changes.

$ ./oarphcli --test-all
Outside the container, run unit tests in all environments.

$ ./oarphcli --build-env --push-as-latest
Rebuild all dockerized environments


## Release Workflow

To update the oarphpy package version, first edit `oarphpy/__init__.py`.
Push your changes to the master branch, then on master run:
$ ./oarphcli --release

"""

import os
import subprocess
import sys

## Logging
import logging
LOG_FORMAT = "%(asctime)s\t%(name)-4s %(process)d : %(message)s"
log = logging.getLogger("op")
log.setLevel(logging.INFO)
console_handler = logging.StreamHandler(sys.stderr)
console_handler.setFormatter(logging.Formatter(LOG_FORMAT))
log.addHandler(console_handler)

# The repository name, e.g. 'oarphpy' in the image name 'oarphpy/full:latest' 
# for the public repository. Toggle this param to use a private repository 
# (e.g. in a private docker registry).
OP_DOCKER_REPOSITORY = os.environ.get('OP_DOCKER_REPOSITORY', 'oarphpy')

# Name of container for the dockerized dev environment
OP_CONTAINER_NAME = os.environ.get('OP_CONTAINER_NAME', 'oarphpy')


## Utils 

def get_oarphpy_version(op_root):
  path = os.path.join(op_root, 'oarphpy/__init__.py')
  
  try:
    # System python is 3.x
    import importlib.machinery
    m = importlib.machinery.SourceFileLoader('_', path).load_module()
  except Exception as e:
    # System python is 2.x; this is deprecated in 3.x
    import imp
    m = imp.load_source('_', path)

  return m.__version__


def run_cmd(cmd):
  cmd = cmd.replace('\n', '').strip()
  log.info("Running %s ..." % cmd)
  subprocess.check_call(cmd, shell=True)
  log.info("... done with %s " % cmd)


def build_and_push_3p_images():
  """Build and push third-party docker images.  This process does not need
  to be run often but is provided here for self-documentation and 
  maintenance releases (e.g. nVidia driver upgrades)
  """

  log.info(
    "Building Lambda Stack docker images (these are not on DockerHub). "
    "FMI see https://github.com/lambdal/lambda-stack-dockerfiles")
  run_cmd("""
    cd /tmp &&
    (rm -rf lambda-stack-dockerfiles || true) &&
    git clone https://github.com/lambdal/lambda-stack-dockerfiles &&
    cd lambda-stack-dockerfiles &&
    git checkout d762400d61636c074533416674426a84cc4d8992 &&
    docker build -t oarphpy/lambda-stack:22.04 -f Dockerfile.jammy . &&
    docker push oarphpy/lambda-stack:22.04
    """)


class DockerEnv(object):
  """Handle for a single Dockerized environment"""

  DOCKERFILE_PATH = None
  IMAGE_NAME = None
  IMAGE_VERSION = None
  SRC_ROOT = None
  DOCKER_REPOSITORY = OP_DOCKER_REPOSITORY

  @staticmethod
  def get_all_envs(op_root, op_version=None):
    """Spider the `docker/` directory of oarphpy and return a handle
    based on each Dockerfile."""

    if not op_version:
      op_version = get_oarphpy_version(op_root)
    
    envs = []
    dockers_dir = os.path.join(op_root, 'docker')
    for fname in os.listdir(dockers_dir):
      if fname.endswith('.Dockerfile'):
        class Env(DockerEnv):
          DOCKERFILE_PATH = os.path.join(dockers_dir, fname)
          IMAGE_NAME = fname.replace('.Dockerfile', '')
          IMAGE_VERSION = op_version
          SRC_ROOT = op_root
        envs.append(Env)
    return envs

  @classmethod
  def full_image_name(cls):
    return "%s/%s:%s" % (
      cls.DOCKER_REPOSITORY , cls.IMAGE_NAME, cls.IMAGE_VERSION)

  @classmethod
  def build(cls, and_push=True):
    image = cls.full_image_name()
    CMD = """
      docker build -t {image} -f {dockerfile} {rootdir}
    """.format(
      image=image,
      dockerfile=cls.DOCKERFILE_PATH,
      rootdir=cls.SRC_ROOT)
    run_cmd(CMD)
    if and_push:
      run_cmd('docker push ' + image)

  @classmethod
  def push_as_latest(cls):
    image = cls.full_image_name()
    latest = image.split(':')[0]
    run_cmd('docker tag ' + image + ' ' + latest)
    run_cmd('docker push ' + latest)

  @classmethod
  def start(cls, container_name=OP_CONTAINER_NAME, mnt_local_root=True):
    image = cls.full_image_name()
    
    # TODO: Features to consider bringing back:
    # * --env-file for GCS / AWS keys
    # * IVY2_PERSISTED_DIR for spark packages
    # * host-persisted temp dir
    # * '--cap-add=SYS_PTRACE --security-opt seccomp=unconfined'
    #    for valgrind / gdb
    # * -v /:/outer_root

    local_mount = ''
    if mnt_local_root:
      local_mount = '-v `pwd`:/opt/oarphpy:z'

    CMD = """
      docker run
        --name {container_name}
        -d -it -P --net=host
        {local_mount}
          {docker_image} sleep infinity || docker start {container_name} || true
    """.format(
          container_name=container_name,
          local_mount=local_mount,
          docker_image=image)
    run_cmd(CMD)

  @classmethod
  def shell(cls, container_name=OP_CONTAINER_NAME):
    cls.start(container_name=container_name)
    EXEC_CMD = 'docker exec -it %s bash' % container_name
    os.execvp("docker", EXEC_CMD.split(' '))

  @classmethod
  def remove(cls, container_name=OP_CONTAINER_NAME):
    try:
      run_cmd('docker rm -f %s' % container_name)
    except Exception:
      pass
    log.info("Removed container %s" % container_name)

  @classmethod
  def run_cmd(
          cls,
          cmd,
          container_name=None,
          force_build=False,
          mnt_local_root=True,
          rm=True):
    """Run `cmd` in a container, and potentially build the needed docker image
    if it doesn't exist."""
    image = cls.full_image_name()
    
    have_image = False
    if not force_build:
      try:
        run_cmd('docker image inspect %s > /dev/null' % image)
        have_image = True
      except Exception:
        pass

    if not have_image:
      log.info("Don't have %s, trying to build ..." % image)
      cls.build(and_push=False)
      log.info("... done building.")

    log.info("Using docker image %s" % image)


    ### Run `cmd`!
    if not container_name:
      container_name = 'oarphpy-temp'

    cls.start(container_name=container_name, mnt_local_root=mnt_local_root)
    RUN_CMD = 'docker exec -it %s %s' % (container_name, cmd)
    run_cmd(RUN_CMD)

    if rm:
      cls.remove(container_name=container_name)

  @classmethod
  def run_tests(cls):
    # Always use a clean run
    cls.remove(container_name='op-test')
    
    # Test!
    if 'py2' in cls.IMAGE_NAME:
      CMD = 'python setup.py test'
    else:
      CMD = 'python3 setup.py test'
    cls.run_cmd(
      CMD,
      container_name='op-test',
      mnt_local_root=False,
      force_build=True,
      rm=True)
  
  @classmethod
  def test_notebooks(cls, op_root):
    # NB: notebooks are designed to run in the [full] environment
    notebooks_dir = os.path.join(op_root, 'notebooks')
    for fname in os.listdir(notebooks_dir):
      if fname.endswith('.ipynb'):
        log.info("Testing notebook %s ..." % fname)
        # For helpful discussion, see http://www.blog.pythonlibrary.org/2018/10/16/testing-jupyter-notebooks/
        CMD = """
          jupyter-nbconvert \
            --ExecutePreprocessor.timeout=3600 \
              --to notebook --execute --output /tmp/out \
                %s
        """ % os.path.join('/opt/oarphpy/notebooks', fname)
        cls.run_cmd(
          CMD,
          container_name='op-notebook-test',
          mnt_local_root=False,
          force_build=True,
          rm=True)
        log.info("... done testing notebook %s ." % fname)


def create_arg_parser():
  import argparse
  
  parser = argparse.ArgumentParser(
                      description=DESC,
                      formatter_class=argparse.RawDescriptionHelpFormatter)
  
  # Configuration
  parser.add_argument(
    '--root', default=os.path.dirname(os.path.abspath(__file__)),
    help='Use source at this root directory [default %(default)s]')
  parser.add_argument(
    '--devenv', default='full',
    help=('Use this environment as the dev environment for --shell '
          '[default %(default)s]'))

  # Actions
  parser.add_argument(
    '--shell', default=False, action='store_true',
    help='Drop into a dockerized shell w/ the given `--devenv`')
  parser.add_argument(
    '--shell-rm', default=False, action='store_true',
    help='Remove the oarphpy dev env container')
  
  parser.add_argument(
    '--build-env', default=False, action='store_true',
    help='Build all oarphpy docker images')
  parser.add_argument(
    '--build-and-push-3p-images', default=False, action='store_true',
    help='Build and push third party images')
  parser.add_argument(
    '--push-as-latest', default=False, action='store_true',
    help='Tag Docker images at latest and push them')
  
  parser.add_argument(
    '--test', default='',
    help='Run unit tests in the given oarphpy docker environment')
  parser.add_argument(
    '--test-all', default=False, action='store_true',
    help='Run unit tests in all oparhpy docker environments')
  parser.add_argument(
    '--test-notebooks', default=False, action='store_true',
    help='Smoke-test the Jupyter notebooks in the `full` environment')

  parser.add_argument(
    '--release', default=False, action='store_true',
    help='Run the release workflow')

  return parser


def main(args=None):
  if not args:
    parser = create_arg_parser()
    args = parser.parse_args()
  
  if args.build_env or args.push_as_latest or args.test_all or args.test:
    for env in DockerEnv.get_all_envs(args.root):
      if args.build_env:
        env.build(and_push=False)
      if args.test_all:
        env.run_tests()
      elif args.test:
        if env.IMAGE_NAME == args.test:
          env.run_tests()
        
        # Print a helful error message at some point
        CHOICES = [ee.IMAGE_NAME for ee in DockerEnv.get_all_envs(args.root)]
        assert args.test in CHOICES, "Invalid choice %s, choices %s" % (
          args.test, CHOICES)

      if args.push_as_latest:
        env.push_as_latest()

  elif args.build_and_push_3p_images:
    build_and_push_3p_images()

  elif args.shell or args.shell_rm:
    devenv = None
    for env in DockerEnv.get_all_envs(args.root):
      if env.IMAGE_NAME == args.devenv:
        devenv = env
    assert devenv, "Could not find env %s" % args.devenv

    if args.shell:
      devenv.shell()
    elif args.shell_rm:
      devenv.remove()
  
  elif args.test_notebooks:
    fullenv = None
    for env in DockerEnv.get_all_envs(args.root):
      if 'full' in env.IMAGE_NAME:
        fullenv = env
    assert fullenv, "Could not find full env"
    fullenv.test_notebooks(args.root)
  
  elif args.release:
    version = get_oarphpy_version(args.root)
    fullenv = None
    py2env = None
    for env in DockerEnv.get_all_envs(args.root):
      if 'full' in env.IMAGE_NAME:
        fullenv = env
      elif 'py2' in env.IMAGE_NAME:
        py2env = env
    assert fullenv, "Could not find full env"
    assert py2env, "Could not find py2 env"

    def log_msg(msg):
      for _ in range(3):
        log.info("=== Release %s: %s ===" % (version, msg))
    
    log_msg("Start tests ...")
    for env in DockerEnv.get_all_envs(args.root):
      env.run_tests()
    log_msg("Test success!")

    log_msg("Pushing Docker Environments ...")
    for env in DockerEnv.get_all_envs(args.root):
      env.build()
      env.push_as_latest()
    log_msg("Push success!")

    log_msg("Generating archive for PyPI ...")
    py2env.run_cmd("python setup.py sdist bdist_egg bdist_wheel")
    fullenv.run_cmd("python3 setup.py sdist bdist_egg bdist_wheel")
    run_cmd("twine check dist/*")
    log_msg("Generation success!")

    # NB: Careful! PyPI does not allow file name re-use!
    # https://test.pypi.org/help/#file-name-reuse
    # Workaround: https://www.python.org/dev/peps/pep-0440/#post-releases
    log_msg("Sending to TEST PyPI ...")
    run_cmd("""
      python -m twine upload --repository-url https://test.pypi.org/legacy/ dist/*
    """)
    log_msg("Sent!")

    log_msg("Sending to REAL PyPI ...")
    run_cmd("twine upload dist/*")
    log_msg("Sent!")

    log_msg("Rebuilding master docs via sphinx ...")
    run_cmd("git checkout master")
    fullenv.run_cmd("bash -c 'cd docs && make html'")
    CMD = """
      git checkout gh-pages &&
      cp -v -r docs/build/html/* ./ &&
      git commit -am 'Docs update from oarphcli --release' &&
      git push origin gh-pages
    """
    run_cmd(CMD)
    log_msg("Docs built and pushed!")

if __name__ == '__main__':
  main()