Configuring the NIM at Runtime

The AlphaFold2 NIM allows a number of customizations at runtime that allow

• Tuning the performance of the NIM

• Taking advantage of your unique hardware

• Modifying runtime variables (such as the location of the NIM cache)

Below, we detail the various AlphaFold2 NIM-specific runtime configuration options.

Let’s first review how to start the NIM:

            

            
export LOCAL_NIM_CACHE=~/.cache/nim

docker run --rm --name alphafold2 --runtime=nvidia \
    -e CUDA_VISIBLE_DEVICES=0 \
    -e NGC_CLI_API_KEY \
    -v $LOCAL_NIM_CACHE:/opt/nim/.cache \
    -p 8000:8000 \
    nvcr.io/nim/deepmind/alphafold2:1.0.0
        

Remember:

• the -p option sets the port for the NIM

• The -e options defined environment variables, which are passed into the NIM’s container at runtime.

• --rm removes the container when it exists

• -it allows interacting with the container directly at the CLI

The location of the NIM cache can be changed on both the local file system and within the NIM. This is useful when a larger drive is needed to store the model weights, for example.

To change the NIM cache location on the host file system, change the source mount in the above command, which we set in the environment variable LOCAL_NIM_CACHE:

            

            
export LOCAL_NIM_CACHE=/mount/largedisk/nim/.cache

docker run --rm --name alphafold2 --runtime=nvidia \
    -e CUDA_VISIBLE_DEVICES=0 \
    -e NGC_CLI_API_KEY \
    -v $LOCAL_NIM_CACHE:/opt/nim/.cache \
    -p 8000:8000 \
    nvcr.io/nim/deepmind/alphafold2:1.0.0
        

This will only affect the location outside the container. To change the location of the cache inside the container, you must:

1. Make sure to mount the new destination path

2. Provide the new path in the NIM_CACHE_PATH environment variable.

Here’s an example where we move the cache within the container to /workdir/cache/af2/:

            

            
export LOCAL_NIM_CACHE=/mount/largedisk/nim/.cache

docker run --rm --name alphafold2 --runtime=nvidia \
    -e CUDA_VISIBLE_DEVICES=0 \
    -e NGC_CLI_API_KEY \
    -e NIM_CACHE_PATH=/workdir/cache/af2 \## Note: we must set this environment variable...
    -v $LOCAL_NIM_CACHE:/workdir/cache/af2 \## and provide the destination mount.
    -p 8000:8000 \
    nvcr.io/nim/deepmind/alphafold2:1.0.0
        

If you have other HTTP or gRPC servers running (for example, other NIMs), you may need to make the 8000 port available by using another port for your NIM. This can be done by:

1. Changing the exposed port by setting the -p option.

2. Setting the NIM_HTTP_API_PORT environment variable to the new port.

Here’s an example where we set our NIM to run on port 7979:

            

            
export LOCAL_NIM_CACHE=/mount/largedisk/nim/.cache

docker run --rm --name alphafold2 --runtime=nvidia \
    -e CUDA_VISIBLE_DEVICES=0 \
    -e NGC_CLI_API_KEY \
    -e NIM_HTTP_API_PORT=7979 \## We must set the NIM_HTTP_API_PORT environment variable...
    -e NIM_CACHE_PATH=/workdir/cache/af2 \
    -v $LOCAL_NIM_CACHE:/workdir/cache/af2 \
    -p 7979:7979 \## as well as forward the port to host.
    nvcr.io/nim/deepmind/alphafold2:1.0.0
        

The AlphaFold2 NIM has been optimized in many ways. One such optimization includes the ability to tune the MSA process for better performance on your machine, based upon the number of available CPU cores, number of available GPUs, and speed of your available SSD (note: we do not recommend running the AlphaFold2 NIM on a machine that uses a hard disk drive for storage, as it will significantly decrease the performance of the NIM).

There are two layers of tunable parallelism for MSA within the AlphaFold2 NIM, set by two different environment variables:

1. NIM_PARALLEL_MSA_RUNNERS: this variable controls the coarse-grained parallelism of the MSA process.

2. NIM_PARALLEL_THREADS_PER_MSA: this variable controls the number of threads used by each MSA process.

            

            
NIM_PARALLEL_MSA_RUNNERS=4
NIM_PARALLEL_THREADS_PER_MSA=12
        

To modify these variables at runtime, pass them as environment variables in the NIM startup command. For example, on a 64-core machine, when running many MSA databases, we might increase NIM_PARALLEL_MSA_RUNNERS to 6 and reduce NIM_PARALLEL_THREADS_PER_MSA to 10 so that we fit our workloads to our machine well:

            

            
docker run --rm --name alphafold2 --runtime=nvidia \
    -e CUDA_VISIBLE_DEVICES=0 \
    -e NGC_CLI_API_KEY \
    -e NIM_PARALLEL_MSA_RUNNERS=6 \## We set the number of parallel runners to 6, perhaps because we have used many MSA databases...
    -e NIM_PARALLEL_THREADS_PER_MSA=10 \## and we set the number of threads per MSA to 12, to 
    -v $LOCAL_NIM_CACHE:/opt/nim/.cache \
    -p 8000:8000 \
    nvcr.io/nim/deepmind/alphafold2:1.0.0