Skip to content

Commit

Permalink
feat: support container run type enroot. (determined-ai#552)
Browse files Browse the repository at this point in the history 
* feat: support container run type enroot.

* feat: support container run type enroot.
  • Loading branch information
@CanmingCobble @eecsliu
CanmingCobble authored and eecsliu committed Dec 1, 2022
1 parent bfab490 commit ec2903a
Show file tree
Hide file tree
Showing 6 changed files with 153 additions and 122 deletions.
7 changes: 5 additions & 2 deletions master/internal/config/dispatcher_resource_manager_config.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import (
const (
singularity = "singularity"
podman = "podman"
enroot = "enroot"
)

// DispatcherResourceManagerConfig is the object that stores the values of
Expand Down Expand Up @@ -48,8 +49,10 @@ type DispatcherSecurityConfig struct {

// Validate performs validation.
func (c DispatcherResourceManagerConfig) Validate() []error {
// Allowed values for the container run type are either 'singularity' or 'podman'
if !(c.LauncherContainerRunType == singularity || c.LauncherContainerRunType == podman) {
// Allowed values for the container run type are either 'singularity', 'podman' or 'enroot'
if !(c.LauncherContainerRunType == singularity ||
c.LauncherContainerRunType == podman ||
c.LauncherContainerRunType == enroot) {
return []error{fmt.Errorf("invalid launch container run type: '%s'", c.LauncherContainerRunType)}
}
return nil
Expand Down
5 changes: 5 additions & 0 deletions master/internal/config/dispatcher_resource_manager_config_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,11 @@ func TestDispatcherResourceManagerConfig_Validate(t *testing.T) {
fields: fields{LauncherContainerRunType: "podman"},
want: nil,
},
{
name: "enroot case",
fields: fields{LauncherContainerRunType: "enroot"},
want: nil,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
Expand Down
39 changes: 22 additions & 17 deletions master/pkg/tasks/dispatcher_task.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,10 @@ import (
)

const (
trueValue = "true"
podman = "podman"
trueValue = "true"
singularity = "singularity"
podman = "podman"
enroot = "enroot"
// dispatcherEntrypointScriptResource is the script to handle container initialization
// before transferring to the defined entrypoint script.
dispatcherEntrypointScriptResource = "dispatcher-wrapper.sh"
Expand All @@ -34,6 +36,7 @@ const (
containerTmpDeterminedDir = "/determined/"
singularityCarrierSlurm = "com.cray.analytics.capsules.carriers.hpc.slurm.SingularityOverSlurm"
podmanCarrierSlurm = "com.cray.analytics.capsules.carriers.hpc.slurm.PodmanOverSlurm"
enrootCarrierSlurm = "com.cray.analytics.capsules.carriers.hpc.slurm.EnrootOverSlurm"
singularityCarrierPbs = "com.cray.analytics.capsules.carriers.hpc.pbs.SingularityOverPbs"
podmanCarrierPbs = "com.cray.analytics.capsules.carriers.hpc.pbs.PodmanOverPbs"
unspecifiedSlotsPerNode = 0
Expand Down Expand Up @@ -117,18 +120,20 @@ func (t *TaskSpec) ToDispatcherManifest(
payload.SetId("com.cray.analytics.capsules.generic.container")
payload.SetVersion("latest")

if containerRunType == podman {
if isPbsLauncher {
payload.SetCarriers([]string{podmanCarrierPbs})
} else {
payload.SetCarriers([]string{podmanCarrierSlurm})
}
} else {
if isPbsLauncher {
payload.SetCarriers([]string{singularityCarrierPbs})
} else {
payload.SetCarriers([]string{singularityCarrierSlurm})
}
// will add case for enroot over pbs
switch {
case isPbsLauncher && containerRunType == podman:
payload.SetCarriers([]string{podmanCarrierPbs})
case !isPbsLauncher && containerRunType == podman:
payload.SetCarriers([]string{podmanCarrierSlurm})
case isPbsLauncher && containerRunType == singularity:
payload.SetCarriers([]string{singularityCarrierPbs})
case !isPbsLauncher && containerRunType == singularity:
payload.SetCarriers([]string{singularityCarrierSlurm})
case containerRunType == enroot:
payload.SetCarriers([]string{enrootCarrierSlurm})
default:
payload.SetCarriers([]string{singularityCarrierSlurm})
}

// Create payload launch parameters
Expand Down Expand Up @@ -219,7 +224,7 @@ func (t *TaskSpec) ToDispatcherManifest(
launchParameters.SetData(mounts)

envVars, err := getEnvVarsForLauncherManifest(
t, masterHost, masterPort, certificateName, userWantsDirMountedOnTmp, slotType)
t, masterHost, masterPort, certificateName, userWantsDirMountedOnTmp, slotType, containerRunType)
if err != nil {
return nil, "", "", err
}
Expand Down Expand Up @@ -450,7 +455,7 @@ func encodeArchiveParameters(
// Gets the environment variables that are to be added to the Launcher's manifest.
func getEnvVarsForLauncherManifest(
taskSpec *TaskSpec, masterHost string, masterPort int, certificateName string,
tmpMount bool, slotType device.Type,
tmpMount bool, slotType device.Type, containerRunType string,
) (map[string]string, error) {
// Hash map containing the environment variables.
m := make(map[string]string)
Expand Down Expand Up @@ -520,7 +525,7 @@ func getEnvVarsForLauncherManifest(

// If the user has not configured a bind mount of /tmp trigger
// dispatcher-wrapper.sh to make it local to the container.
if !tmpMount {
if !tmpMount && containerRunType != enroot {
m["DET_CONTAINER_LOCAL_TMP"] = "1"
}

Expand Down
15 changes: 13 additions & 2 deletions master/pkg/tasks/dispatcher_task_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -743,7 +743,7 @@ func Test_getEnvVarsForLauncherManifest(t *testing.T) {
}

envVars, err := getEnvVarsForLauncherManifest(ts,
"masterHost", 8888, "certName", false, device.CUDA)
"masterHost", 8888, "certName", false, device.CUDA, "singularity")

assert.NilError(t, err)
assert.Assert(t, len(envVars) > 0)
Expand All @@ -766,6 +766,16 @@ func Test_getEnvVarsForLauncherManifest(t *testing.T) {

assert.Equal(t, envVars["cpu"], "default")
assert.Equal(t, envVars["myenv"], "xyz")

envVarsPodman, _ := getEnvVarsForLauncherManifest(ts,
"masterHost", 8888, "certName", false, device.CUDA, "podman")
assert.Equal(t, envVarsPodman["DET_CONTAINER_LOCAL_TMP"], "1")

// test DET_CONTAINER_LOCAL_TMP is not in ENVs
envVarsEnroot, _ := getEnvVarsForLauncherManifest(ts,
"masterHost", 8888, "certName", false, device.CUDA, "enroot")
_, ok := envVarsEnroot["DET_CONTAINER_LOCAL_TMP"]
assert.Equal(t, ok, false)
}

func Test_getEnvVarsForLauncherManifestErr(t *testing.T) {
Expand All @@ -787,7 +797,8 @@ func Test_getEnvVarsForLauncherManifestErr(t *testing.T) {
Environment: environment,
}

_, err := getEnvVarsForLauncherManifest(ts, "masterHost", 8888, "certName", false, device.CUDA)
_, err := getEnvVarsForLauncherManifest(ts, "masterHost", 8888, "certName", false,
device.CUDA, "podman")
assert.ErrorContains(t, err, "invalid user-defined environment variable 'cpudefault'")
}

Expand Down
38 changes: 20 additions & 18 deletions master/static/srv/dispatcher-wrapper.sh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@
set -eE
trap 'echo >&2 "FATAL: Unexpected error terminated dispatcher-wrapper container initialization. See error messages above."' ERR


# Controls debug logging for this method
DEBUG=0

Expand All @@ -38,11 +37,10 @@ unset -f $(declare -Ffx | cut -f 3 -d ' ')
# Args: {Level} {Message}...
log_debug() {
if [ $DEBUG == 1 ]; then
echo -e "$*" >&2
echo -e "$*" >&2
fi
}


# Unconditional log method
# Args: {Level} {Message}...
log() {
Expand All @@ -64,25 +62,24 @@ log() {
# TODO: Need to handle Multi Instance GPU (MIG) Format.
# Refer: https://docs.nvidia.com/datacenter/tesla/pdf/NVIDIA_MIG_User_Guide.pdf Section 9.6.1 for
# further information on MIG format
convert_to_gpu_numbers(){
convert_to_gpu_numbers() {
# Process the value of CUDA_VISIBLE_DEVICES and store the values in an array.
# IFS flag is set to "," to process the string as a comma separated list.
IFS=',' read -r -a cuda_device_ids <<< "$CUDA_VISIBLE_DEVICES"
IFS=',' read -r -a cuda_device_ids <<<"$CUDA_VISIBLE_DEVICES"
# Check if the first element is a number.
if [[ "${cuda_device_ids[0]}" =~ ^[[:digit:]]+$ ]]; then
if [[ ${cuda_device_ids[0]} =~ ^[[:digit:]]+$ ]]; then
# Return the value as it is, if it is already in simple number format.
echo "${CUDA_VISIBLE_DEVICES}"
else
# Update the value of CUDA_VISIBLE_DEVICES.
cuda_devices_string=""
error_flag=0
# Below for loop will creates a string in the format "0,1,2,"
for gpu_id in "${cuda_device_ids[@]}"
do
for gpu_id in "${cuda_device_ids[@]}"; do
# Retrieve gpu id in the simple number format using the nvidia-smi command.
simple_gpu_id=$(nvidia-smi --query-gpu=index --format=csv,noheader -i ${gpu_id})
# If the command failed log warning and return the existing value as it is.
if [ $? -ne 0 ];then
if [ $? -ne 0 ]; then
log "ERROR: Failed to retrieve index for GID ${gpu_id} using nvidia-smi." 1>&2
error_flag=1
break
Expand All @@ -104,6 +101,11 @@ convert_to_gpu_numbers(){
# With --writable-tmpfs option / is writable by the user
# and private to the container instance.
LOCALTMP=/
if [ ! -w $LOCALTMP ]; then
# If / is not writable (Enroot), then /tmp is
# container-private and we can just use that.
LOCALTMP=/tmp
fi
# Source volume of all archives to be cloned
ROOT="/determined_local_fs"
# Base of the per-proc copy of tree
Expand All @@ -112,10 +114,10 @@ PROCDIR_ROOT="$ROOT/procs"
PROCDIR="$PROCDIR_ROOT/$SLURM_PROCID"

# Create clone of any directories under /dispatcher for this process and setup links
if [ -d $ROOT/run ] ; then
if [ -d $ROOT/run ]; then
mkdir -p $PROCDIR
for dir in $ROOT/*; do
if [[ -d $dir && $dir != $PROCDIR_ROOT ]] ; then
if [[ -d $dir && $dir != $PROCDIR_ROOT ]]; then
log_debug "INFO: Clone $dir -> $PROCDIR"
cp -p -R $dir $PROCDIR >&2
fi
Expand Down Expand Up @@ -143,7 +145,7 @@ if [ -d $ROOT/run ] ; then
fi

# Localize /tmp as a private folder in the container, if requested.
if [ "$DET_CONTAINER_LOCAL_TMP" == "1" ]; then
if [ "$DET_CONTAINER_LOCAL_TMP" == "1" ]; then
# Create a per-container tmp
mkdir -p $PROCDIR/tmp
# Replace /tmp with a link to our private
Expand Down Expand Up @@ -172,13 +174,13 @@ if [ "$DET_RESOURCES_TYPE" == "slurm-job" ]; then

if [ ! -z "$CUDA_VISIBLE_DEVICES" ]; then
# Test if "nvidia-smi" exists in the PATH before trying to invoking it.
if type nvidia-smi > /dev/null 2>&1 ; then
if type nvidia-smi >/dev/null 2>&1; then
# For Nvidia GPUS, the slot IDs are the device index. Replace the
# newline characters with commas and enclose in square brackets.
# But only include GPUS that are in the CUDA_VISIBLE_DEVICES=0,1,...
VISIBLE_SLOTS="$(nvidia-smi --query-gpu=index --format=csv,noheader | sed -z 's/\n/,/g;s/,$/\n/')"
for device in ${CUDA_VISIBLE_DEVICES//,/ } ; do
if [[ ! "$VISIBLE_SLOTS" == *"$device"* ]]; then
for device in ${CUDA_VISIBLE_DEVICES//,/ }; do
if [[ $VISIBLE_SLOTS != *"$device"* ]]; then
log "WARNING: nvidia-smi reports visible CUDA devices as ${VISIBLE_SLOTS} but does not contain ${device}. May be unable to perform CUDA operations." 1>&2
fi
done
Expand Down Expand Up @@ -208,15 +210,15 @@ if [ "$DET_RESOURCES_TYPE" == "slurm-job" ]; then
if [[ -x /usr/bin/rocm-smi ]]; then
if grep -s /usr/libexec/platform-python /usr/bin/rocm-smi; then
mkdir -p /run/determined/pythonuserbase/bin/
echo -e '#!/bin/bash\npython3 /usr/bin/rocm-smi $*' > /run/determined/pythonuserbase/bin/rocm-smi
echo -e '#!/bin/bash\npython3 /usr/bin/rocm-smi $*' >/run/determined/pythonuserbase/bin/rocm-smi
chmod +x /run/determined/pythonuserbase/bin/rocm-smi
log "INFO: Adding rocm-smi wrapper script /run/determined/pythonuserbase/bin/rocm-smi." 1>&2
fi
fi

if [ ! -z "$ROCR_VISIBLE_DEVICES" ]; then
# Test if "rocm-smi" exists in the PATH before trying to invoking it.
if [ ! type rocm-smi > /dev/null 2>&1 ]; then
if [ ! type rocm-smi ] >/dev/null 2>&1; then
log "WARNING: rocm-smi not found. May be unable to perform ROCM operations." 1>&2
fi
else
Expand Down Expand Up @@ -253,7 +255,7 @@ fi
# result in:
# cat /run/determined/etc/passwd
# username:x:0:0::/run/determined/workdir:/bin/bash
if [ $(whoami) == "root" ] && [ -r /run/determined/etc/passwd ]; then
if [ $(whoami) == "root" ] && [ -r /run/determined/etc/passwd ]; then
log_debug "DEBUG: Running as root inside container, changing agent user passwd entry to uid/gid 0/0."
sed -i "s/\([a-zA-Z0-9]\+\):x:[0-9]\+:[0-9]\+:/\1:x:0:0:/" /run/determined/etc/passwd
fi
Expand Down
Loading

0 comments on commit ec2903a

Please sign in to comment.