conda_queue_env_inline_improved_caching.yaml

specificationVersion: "environment-2023-09"
parameterDefinitions:
# CondaPackages and CondaChannels are compatible with the other Conda queue environment templates
- name: CondaPackages
  type: STRING
  description: >
    This is a space-separated list of Conda package match specifications to
    install for the job. E.g. "blender=4.3" for a job that renders frames in
    Blender 4.3.
  default: ""
  userInterface:
    control: LINE_EDIT
    label: Conda Packages
- name: CondaChannels
  type: STRING
  description: >
    This is a space-separated list of Conda channels from which to install
    packages. Deadline Cloud SMF packages are installed from the
    "deadline-cloud" channel that is configured by Deadline Cloud.

    Add "conda-forge" to get packages from the https://conda-forge.org/
    community, and "defaults" to get packages from Anaconda Inc (make sure
    your usage complies with https://legal.anaconda.com/policies).
  default: "deadline-cloud"
  userInterface:
    control: LINE_EDIT
    label: Conda Channels

- name: NamedCondaEnv
  description: |
    If empty, the Conda environment is created only for the session. If provided,
    a named Conda environment is reused across any jobs/sessions that run on the worker host.
    When reusing an environment, the parameter NamedCondaEnvUpdateAfterMinutes controls
    how long it will use the environment without checking for package updates.

    If the name is "AUTOMATIC", the environment name is determined by taking the hash of the
    CondaPackages and CondaChannels parameters, so that jobs providing identical values for
    these parameters will share the same named Conda environment. Automatically-named environments
    are deleted by the onExit action if they haven't been updated in more than 96 hours.

    If you select a name explicitly, take care to use the environment names consistently, because
    different jobs can define the same named environment in a different way.

    When creating or updating this environment, it prints information about the operations
    to the log file "$CONDA_PREFIX/var/log/conda_queue_env.log".
  type: STRING
  default: "AUTOMATIC"
  userInterface:
    control: LINE_EDIT
    label: Named Conda Environment

- name: NamedCondaEnvAction
  description: |
    When a NamedCondaEnv is selected, controls how to treat it.
      1. ACTIVATE activates the existing named environment, or creates one if it
         doesn't exist. If it's been NamedCondaEnvUpdateAfterMinutes since creation
         or the last update, it updates the environment.
      2. REMOVE_AND_CREATE always removes the existing named environment, and creates
         a new one from scratch.
  type: STRING
  default: "ACTIVATE"
  allowedValues: ["ACTIVATE", "REMOVE_AND_CREATE"]
  userInterface:
    control: DROPDOWN_LIST
    label: Named Conda Environment Action

- name: NamedCondaEnvUpdateAfterMinutes
  description: |
    When NamedCondaEnvAction is ACTIVATE, this controls how long it will use the named
    Conda environment before doing another update.
  type: INT
  default: 600
  userInterface:
    control: SPIN_BOX
    label: Named Conda Environment Update After (Minutes)

- name: RunCondaClean
  description: |
    If set to True, runs the command 'conda clean --yes --all' before creating the Conda environment.
    This removes Conda caches for the package index, package files, and package directories.
  type: STRING
  default: "False"
  allowedValues: ["True", "False"]
  userInterface:
    control: CHECK_BOX
    label: Clean the Conda Cache

environment:
  name: Conda
  script:
    actions:
      onEnter:
        command: "bash"
        args: ["{{Env.File.Enter}}"]
      onExit:
        command: "bash"
        args: ["{{Env.File.Exit}}"]
    embeddedFiles:
    - name: Enter
      filename: conda-queue-env-enter.sh
      type: TEXT
      data: |
        set -euo pipefail

        if [ -z '{{Param.CondaPackages}}' ]; then
            echo "Skipping Conda env as CondaPackages parameter was empty."
            exit 0
        fi

        # Install an error handler to clean the cache if there is an error creating the virtual environment
        function conda_clean_on_error {
            if [ ! "$1" = "0" ]; then
              echo "Error detected, cleaning the Conda cache."
              conda clean --yes --all
            fi
        }
        trap 'conda_clean_on_error $?' EXIT

        function remove_named_env {
            for ENVS_DIR in $(conda info --json | python -c "import json, sys; v = json.load(sys.stdin); print('\n'.join(v['envs_dirs']))"); do
                if [ -d $ENVS_DIR/$1 ]; then
                    echo "Removing directory $ENVS_DIR/$1 for the environment"
                    rm -rf $ENVS_DIR/$1
                    if [ -d $ENVS_DIR/$1 ]; then
                        echo "ERROR: Could not remove the directory. Possibly a permissions error or a process holding a lock."
                        exit 1
                    fi
                fi
            done
        }

        NAMED_CONDA_ENV='{{Param.NamedCondaEnv}}'

        # If the special name 'AUTOMATIC' is provided, use the hash of the CondaChannels and CondaPackages
        # parameters to generate the actual name.
        if [ -n "$NAMED_CONDA_ENV" ] && [ "$NAMED_CONDA_ENV" = 'AUTOMATIC' ]; then
            echo 'Using an automatic name for the Conda environment, based on the hash of these values:'
            echo '  CondaChannels: {{Param.CondaChannels}}'
            echo '  CondaPackages: {{Param.CondaPackages}}'

            NAMED_CONDA_ENV="hashname_$(echo 'Channels/{{Param.CondaChannels}} -- Packages/{{Param.CondaPackages}}' | sha256sum | cut -d " " -f 1 | head -c 24)"
            echo "Automatic name is $NAMED_CONDA_ENV"
        fi

        # If we're not reusing the Conda env, clean up any prior ones
        if [ '{{Param.NamedCondaEnvAction}}' = 'REMOVE_AND_CREATE' ] && [ -n "$NAMED_CONDA_ENV" ]; then
            echo "NamedCondaEnvAction parameter is set to REMOVE_AND_CREATE, removing the virtual environment..."
            remove_named_env "$NAMED_CONDA_ENV"
        fi

        # If requested, clean the Conda package cache
        if [ '{{Param.RunCondaClean}}' = 'True' ]; then
            echo "RunCondaClean parameter is True, cleaning the Conda cache..."
            conda clean --yes --all
        fi

        # Convert the space-separated list of channels into consecutive '-c' channel options
        CHANNEL_OPTS="$(echo '{{Param.CondaChannels}}' | sed -r 's/(\s+|^)(\S)/ -c \2/g')"
        # Put the conda packages list in a variable, as package specs can have characters like '|' in them
        CONDA_PACKAGES='{{Param.CondaPackages}}'

        # Initialize/activate the Conda virtual environment
        if [ -n "$NAMED_CONDA_ENV" ] && conda env list | grep -q "^$NAMED_CONDA_ENV "; then
            echo "Reusing the existing named Conda environment $NAMED_CONDA_ENV."
            echo ""

            # Capture environment variables set by activation
            python '{{Env.File.OpenJDVarsStart}}' .vars
            conda run -n "$NAMED_CONDA_ENV" python '{{Env.File.OpenJDVarsCapture}}' .vars
            ENV_DIR=$(conda run -n "$NAMED_CONDA_ENV" bash -c 'echo $CONDA_PREFIX')

            CUR_TIMESTAMP="$(date +%s)"
            TIMESTAMP_FILE="$ENV_DIR/var/log/conda_queue_env_update_timestamp"
            UPDATE_AFTER='{{Param.NamedCondaEnvUpdateAfterMinutes}}'
            if [ -f "$TIMESTAMP_FILE" ]; then
                PREV_UPDATE_TIMESTAMP="$(cat "$TIMESTAMP_FILE")"
                MINUTES_SINCE_UPDATE="$(python -c "print(($CUR_TIMESTAMP - $PREV_UPDATE_TIMESTAMP) // 60)")"
            else
                echo "The file recording the previous timestamp was not found"
                MINUTES_SINCE_UPDATE="$UPDATE_AFTER"
            fi
            echo "Minutes elapsed since last update of this named Conda environment: $MINUTES_SINCE_UPDATE"

            if [ $MINUTES_SINCE_UPDATE -ge $UPDATE_AFTER ]; then
                echo "Elapsed time greater than or equal to $UPDATE_AFTER minutes, updating the environment packages to the latest"
                conda install -n "$NAMED_CONDA_ENV" --yes --quiet \
                    $CONDA_PACKAGES \
                    $CHANNEL_OPTS

                # Save the channels and packages used in the environment, to help out debugging
                LOGFILE="$ENV_DIR/var/log/conda_queue_env.log"
                echo "Updated $NAMED_CONDA_ENV env at $(date --iso-8601=minutes)" >> "$LOGFILE"
                echo '  CondaChannels: {{Param.CondaChannels}}' >> "$LOGFILE"
                echo '  CondaPackages: {{Param.CondaPackages}}' >> "$LOGFILE"

                # Save a timestamp of when we updated
                echo "$CUR_TIMESTAMP" > "$ENV_DIR/var/log/conda_queue_env_update_timestamp"
            else
                echo "Elapsed time less than $UPDATE_AFTER minutes, skipping updates"
            fi

            echo ""
            conda list -n "$NAMED_CONDA_ENV"
        elif [ -n "$NAMED_CONDA_ENV" ]; then
            echo "Named Conda environment $NAMED_CONDA_ENV not found, creating it."

            # Remove incomplete environments
            remove_named_env "$NAMED_CONDA_ENV"

            conda create --yes --quiet \
                --name "$NAMED_CONDA_ENV" \
                $CONDA_PACKAGES \
                $CHANNEL_OPTS

            # Capture environment variables set by activation
            python '{{Env.File.OpenJDVarsStart}}' .vars
            conda run -n "$NAMED_CONDA_ENV" python '{{Env.File.OpenJDVarsCapture}}' .vars
            ENV_DIR=$(conda run -n $NAMED_CONDA_ENV bash -c 'echo $CONDA_PREFIX')

            # Save the channels and packages used in the environment, to help debugging
            LOGFILE="$ENV_DIR/var/log/conda_queue_env.log"
            mkdir -p "$(dirname $LOGFILE)"
            echo "Created $NAMED_CONDA_ENV env at $(date --iso-8601=minutes)" >> "$LOGFILE"
            echo '  CondaChannels: {{Param.CondaChannels}}' >> "$LOGFILE"
            echo '  CondaPackages: {{Param.CondaPackages}}' >> "$LOGFILE"

            # Save a timestamp of when we updated
            date +%s > "$ENV_DIR/var/log/conda_queue_env_update_timestamp"
        else
            echo "Creating temporary Conda environment in the session directory..."

            ENV_DIR="$(mktemp -d '{{Session.WorkingDirectory}}/.env-XXXXX')"
            conda create --yes --quiet \
                -p "$ENV_DIR" \
                $CONDA_PACKAGES \
                $CHANNEL_OPTS

            # Capture environment variables set by activation
            python '{{Env.File.OpenJDVarsStart}}' .vars
            conda run -p "$ENV_DIR" python '{{Env.File.OpenJDVarsCapture}}' .vars
        fi

        conda run -p "$ENV_DIR" conda info
    - name: Exit
      filename: conda-queue-env-exit.sh
      type: TEXT
      data: |
        set -euo pipefail

        # When using this queue environment on long-lived worker hosts, such as on premises, the
        # automatically named conda environments can accumulate. This code cleans up these environments
        # after they haven't been used for 96 hours, a length of time chosen so that when a farm is idle
        # for a long weekend it doesn't clear all these environment caches.

        # Deactivate the environment so we can delete it if necessary
        set +u
        eval "$(conda shell.bash deactivate)"
        set -u

        ENV_DELETE_AFTER_HOURS=96

        echo "Cleaning up any automatically-named conda environments that weren't updated within $ENV_DELETE_AFTER_HOURS hours."

        AUTO_CONDA_ENVS="$(conda env list | { grep "^hashname_" || true; } | cut -d " " -f 1)"
        REMOVED_ENV=0

        CUR_TIMESTAMP="$(date +%s)"
        for ENV_NAME in $AUTO_CONDA_ENVS; do
            echo "Checking environment $ENV_NAME"
            ENV_DIR=$(conda run -n $ENV_NAME bash -c 'echo $CONDA_PREFIX')

            LOGFILE="$ENV_DIR/var/log/conda_queue_env.log"
            if [ -f "$LOGFILE" ]; then
                # Creation and each update outputs 3 lines to the log, so print the last 3
                tail -3 "$LOGFILE"
            fi

            TIMESTAMP_FILE="$ENV_DIR/var/log/conda_queue_env_update_timestamp"
            if [ -f "$TIMESTAMP_FILE" ]; then
                PREV_UPDATE_TIMESTAMP="$(cat "$TIMESTAMP_FILE")"
            else
                echo "The file recording the previous timestamp was not found"
                PREV_UPDATE_TIMESTAMP=0
            fi
            HOURS_SINCE_UPDATE="$(python -c print"(($CUR_TIMESTAMP - $PREV_UPDATE_TIMESTAMP) // 3600)")"
            echo "Environment was last updated $HOURS_SINCE_UPDATE hours ago."

            if [ $HOURS_SINCE_UPDATE -gt $ENV_DELETE_AFTER_HOURS ]; then
                echo "Elapsed time $HOURS_SINCE_UPDATE is greater than $ENV_DELETE_AFTER_HOURS hours, removing the environment."
                conda env remove --yes -q -n $ENV_NAME
                REMOVED_ENV=1
            fi
        done

        if [ "$REMOVED_ENV" = "1" ]; then
            echo "Cleaning the conda cache to reclaim disk space."
            conda clean --yes --all
        fi

    - name: OpenJDVarsStart
      filename: openjd-vars-start.py
      type: TEXT
      data: |
        import json
        import os
        import sys

        # Exclude the env var "_" as it has special meaning to shells
        before = dict(os.environ)
        if "_" in before:
            del before["_"]

        with open(sys.argv[1], "w", encoding="utf8") as f:
            json.dump(before, f)
    - name: OpenJDVarsCapture
      filename: openjd-vars-capture.py
      type: TEXT
      data: |
        import json
        import os
        import sys

        # Get the snapshot from `openjd-vars-start.py`, and the current environment state.
        with open(sys.argv[1], "r", encoding="utf8") as f:
            before = json.load(f)
        after = dict(os.environ)
        # Exclude the env var "_" as it has special meaning to shells
        if "_" in after:
            del after["_"]

        # Identify the modified and deleted environment variables
        vars_to_put = {k: v for k, v in after.items() if v != before.get(k)}
        vars_to_delete = {k for k in before if k not in after}

        # Print the env var changes following the Open Job Description specification
        for k, v in vars_to_put.items():
            kv = json.dumps(f"{k}={v}", ensure_ascii=True)
            print(f"openjd_env: {kv}")
        for k in vars_to_delete:
            print(f"openjd_unset_env: {k}")