File size: 2,329 Bytes
79deb35
 
cf62cfd
 
 
 
79deb35
2501a37
 
 
 
cf62cfd
 
2501a37
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
cf62cfd
2501a37
 
 
 
 
cf62cfd
 
2501a37
729740d
2501a37
729740d
 
cf62cfd
729740d
cf62cfd
79deb35
ece0211
 
 
 
 
 
 
 
120b809
ece0211
 
e6937e8
 
 
 
 
 
 
79deb35
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
#!/bin/bash

# Export specific ENV variables to /etc/rp_environment
echo "Exporting environment variables..."
printenv | grep -E '^RUNPOD_|^PATH=|^_=' | sed 's/^\(.*\)=\(.*\)$/export \1="\2"/' >> /etc/rp_environment
echo 'source /etc/rp_environment' >> ~/.bashrc

add_keys_to_authorized() {
    local key_value=$1

    # Create the ~/.ssh directory and set permissions
    mkdir -p ~/.ssh
    chmod 700 ~/.ssh

    # Create the authorized_keys file if it doesn't exist
    touch ~/.ssh/authorized_keys

    # Initialize an empty key variable
    local key=""

    # Read the key variable word by word
    for word in $key_value; do
        # Check if the word looks like the start of a key
        if [[ $word == ssh-* ]]; then
            # If there's a key being built, add it to the authorized_keys file
            if [[ -n $key ]]; then
                echo $key >> ~/.ssh/authorized_keys
            fi
            # Start a new key
            key=$word
        else
            # Append the word to the current key
            key="$key $word"
        fi
    done

    # Add the last key to the authorized_keys file
    if [[ -n $key ]]; then
        echo $key >> ~/.ssh/authorized_keys
    fi

    # Set the correct permissions
    chmod 600 ~/.ssh/authorized_keys
    chmod 700 -R ~/.ssh
}

if [[ $PUBLIC_KEY ]]; then
    # runpod
    add_keys_to_authorized "$PUBLIC_KEY"
    # Start the SSH service in the background
    service ssh start
elif [[ $SSH_KEY ]]; then
    # latitude.sh
    add_keys_to_authorized "$SSH_KEY"
    # Start the SSH service in the background
    service ssh start
else
    echo "No PUBLIC_KEY or SSH_KEY environment variable provided, not starting openSSH daemon"
fi

# Check if JUPYTER_PASSWORD is set and not empty
if [ -n "$JUPYTER_PASSWORD" ]; then
    # Set JUPYTER_TOKEN to the value of JUPYTER_PASSWORD
    export JUPYTER_TOKEN="$JUPYTER_PASSWORD"
fi

if [ "$JUPYTER_DISABLE" != "1" ]; then
    # Run Jupyter Lab in the background
    jupyter lab --port=8888 --ip=* --allow-root --ServerApp.allow_origin=* &
fi

if [ ! -d "/workspace/data/axolotl-artifacts" ]; then
    mkdir -p /workspace/data/axolotl-artifacts
fi
if [ ! -L "/workspace/axolotl/outputs" ]; then
    ln -sf /workspace/data/axolotl-artifacts /workspace/axolotl/outputs
fi

# Execute the passed arguments (CMD)
exec "$@"