#!/usr/bin/env -S just --justfile

# https://github.com/casey/just?tab=readme-ov-file#dotenv-settings

set dotenv-load

# https://github.com/casey/just?tab=readme-ov-file#export

set export

# don't fail fast here --> the `setup` command will check this!

COMMIT := `git rev-parse HEAD || true`

IMAGE_TAG := "bionemo2-" + COMMIT

DEV_IMAGE_TAG := "dev-" + IMAGE_TAG

DATE := `date --iso-8601=seconds -u`

LOCAL_ENV := '.env'

DOCKER_REPO_PATH := '/workspace/bionemo2'

LOCAL_REPO_PATH := `realpath $(pwd)`

[private]

default:

  @just --list

###############################################################################

[private]

check_preconditions:

  #!/usr/bin/env bash

  version_ge() {

      # Returns 0 (true) if $1 >= $2, 1 (false) otherwise

      [ "$(printf '%s\n' "$1" "$2" | sort -V | head -n1)" = "$2" ]

  }

  if [[ $(command -v git) ]]; then

    commit=$(git rev-parse HEAD)

    if [[ "$?" != "0" ]]; then

      echo "ERROR: must run from within git repository!"

      exit 1

    fi

  else

    echo "ERROR: git is not installed!"

    exit 1

  fi

  if [[ ! $(command -v docker) ]]; then

    echo "ERROR: docker is not installed!"

    exit 1

  fi

  docker_version=$(docker --version | awk -F'[, ]' '{print $3}')

  required_docker_version='23.0.1'

  if ! version_ge "$docker_version" "$required_docker_version"; then

      echo "Error: Docker version $required_docker_version or higher is required. Current version: $docker_version"

      exit 1

  fi

# Checks for installed programs (docker, git, etc.), their versions, and grabs the latest cache image.

setup: check_preconditions

  ./internal/scripts/setup_env_file.sh

  @echo "Pulling updated cache..."

  docker pull ${IMAGE_REPO}:${CACHE_TAG} || true

[private]

assert_clean_git_repo:

  #!/usr/bin/env bash

  git diff-index --quiet HEAD --

  exit_code="$?"

  if [[ "${exit_code}" == "128" ]]; then

      echo "ERROR: Cannot build image if not in bionemo git repository!"

      exit 1

  elif [[ "${exit_code}" == "1" ]]; then

      echo "ERROR: Repository is dirty! Commit all changes before building image!"

      exit  2

  elif [[ "${exit_code}" == "0" ]]; then

      echo "ok" 2> /dev/null

  else

      echo "ERROR: Unknown exit code for `git diff-index`: ${exit_code}"

      exit 1

  fi

###############################################################################

[private]

build image_tag target: setup assert_clean_git_repo

  DOCKER_BUILDKIT=1 docker buildx build \

  -t ${IMAGE_REPO}:{{image_tag}} \

  --target={{target}} \

  --load \

  --cache-to type=inline \

  --cache-from ${IMAGE_REPO}:${CACHE_TAG} \

  --label com.nvidia.bionemo.git_sha=${COMMIT} \

  --label com.nvidia.bionemo.created_at=${DATE} \

  -f ./Dockerfile \

  .

# Builds the release image.

build-release:

  @just build ${IMAGE_TAG} release

# Builds the development image.

build-dev:

  @just build ${DEV_IMAGE_TAG} development

###############################################################################

[private]

run is_dev is_interactive image_tag cmd: setup

  #!/usr/bin/env bash

  DOCKER_VERSION=$(docker version | grep -i version | head -1 | awk '{print $2}')

  DOCKER_VERSION_WITH_GPU_SUPPORT='19.03.0'

  if [ "$DOCKER_VERSION_WITH_GPU_SUPPORT" == "$(echo -e "$DOCKER_VERSION\n$DOCKER_VERSION_WITH_GPU_SUPPORT" | sort -V | head -1)" ]; then

      PARAM_RUNTIME="--gpus all"

  else

      PARAM_RUNTIME="--runtime=nvidia"

  fi

  docker_cmd="docker run \

  --rm \

  --network host \

  ${PARAM_RUNTIME} \

  -p ${JUPYTER_PORT}:8888 \

  --shm-size=4g \

  -e TMPDIR=/tmp/ \

  -e NUMBA_CACHE_DIR=/tmp/ \

  -e WANDB_API_KEY=$WANDB_API_KEY \

  -e NGC_CLI_API_KEY=$NGC_CLI_API_KEY \

  -e NGC_CLI_ORG=$NGC_CLI_ORG \

  -e NGC_CLI_TEAM=$NGC_CLI_TEAM \

  -e NGC_CLI_FORMAT_TYPE=$NGC_CLI_FORMAT_TYPE \

  -e AWS_ENDPOINT_URL \

  -e AWS_REGION \

  -e AWS_SECRET_ACCESS_KEY \

  -e AWS_ACCESS_KEY_ID \

  -e HOME=${DOCKER_REPO_PATH} \

  -w ${DOCKER_REPO_PATH} \

  -v ${LOCAL_RESULTS_PATH}:${DOCKER_RESULTS_PATH} \

  -v ${LOCAL_DATA_PATH}:${DOCKER_DATA_PATH} \

  -v ${LOCAL_MODELS_PATH}:${DOCKER_MODELS_PATH} \

  -v /etc/passwd:/etc/passwd:ro \

  -v /etc/group:/etc/group:ro \

  -v /etc/shadow:/etc/shadow:ro \

  -v ${HOME}/.ssh:${DOCKER_REPO_PATH}/.ssh:ro \

  -v ${LOCAL_REPO_PATH}/htmlcov:/${DOCKER_REPO_PATH}/htmlcov \

  -u $(id -u):$(id -g)"

  if [[ "{{is_dev}}" == "true" ]]; then

    docker_cmd="${docker_cmd} -v ${LOCAL_REPO_PATH}:${DOCKER_REPO_PATH}"

  fi

  if [[ "{{is_interactive}}" == "true" ]]; then

    docker_cmd="${docker_cmd} -it"

  fi

  docker_cmd="${docker_cmd} ${IMAGE_REPO}:{{image_tag}} {{cmd}}"

  set -xeuo pipefail

  DOCKER_BUILDKIT=1 ${docker_cmd}

[private]

ensure-dev-or-build:

  #!/usr/bin/env bash

  if [[ $(docker images -q "${IMAGE_REPO}:${DEV_IMAGE_TAG}" 2> /dev/null) == "" ]]; then

    echo "Building development image: ${IMAGE_REPO}:${DEV_IMAGE_TAG}"

    just build-dev

  else

    echo "Development image exists:   ${IMAGE_REPO}:${DEV_IMAGE_TAG}"

  fi

# run-dev lets us work with a dirty repository,

# beacuse this is a common state during development

# **AND** we're volume mounting the code, so we'll have the latest state.

# **BUT** we can only do this if we have built an image for DEV_IMAGE_TAG already,

# so we use ensure-dev-or-build which will build the image if it is necessary.

# Image building requires that the git repo state is clean!

#

# Runs an interactive program in the development bionemo image.

run-dev cmd='bash': ensure-dev-or-build

  @just run true true ${DEV_IMAGE_TAG} {{cmd}}

# in contrast, run-release requires a clean repository,

# because users want to know that they're running the **exact** version they expect

# and we're **NOT** volume mounting the code

#

# Runs an interactive program in the release bionemo image.

run-release cmd='bash': build-release assert_clean_git_repo

  @just run false true ${IMAGE_TAG} {{cmd}}

###############################################################################

# Executes pytest in the release image.

test: build-release

  @just run true false ${IMAGE_TAG} 'pytest -v --nbval-lax --cov=bionemo --cov-report term --cov-report=html docs/ scripts/ sub-packages/'