Merge pull request #1372 from clayne/1743815607-wandb-api-key-vs-token

Use correct HF token and WanDB api-key env-vars

Author: bghira

docker-start.sh

@@ -4,7 +4,7 @@
 # This file can then later be sourced in a login shell
 echo "Exporting environment variables..."
 printenv |
-	grep -E '^RUNPOD_|^PATH=|^HF_HOME=|^HUGGING_FACE_HUB_TOKEN=|^_=' |
+	grep -E '^RUNPOD_|^PATH=|^HF_HOME=|^HF_TOKEN=|^HUGGING_FACE_HUB_TOKEN=|^WANDB_API_KEY=|^WANDB_TOKEN=|^_=' |
 	sed 's/^\(.*\)=\(.*\)$/export \1="\2"/' >>/etc/rp_environment
 
 # Add it to Bash login script
@@ -26,9 +26,19 @@ fi
 # Start SSH server
 service ssh start
 
-# Load HF, WanDB tokens
-if [ -n "$HUGGING_FACE_HUB_TOKEN" ]; then huggingface-cli login --token "$HUGGING_FACE_HUB_TOKEN" --add-to-git-credential; else echo "HUGGING_FACE_HUB_TOKEN not set; skipping login"; fi
-if [ -n "$WANDB_TOKEN" ]; then wandb login "$WANDB_TOKEN"; else echo "WANDB_TOKEN not set; skipping login"; fi
+# Login to HF
+if [[ -n "${HF_TOKEN:-$HUGGING_FACE_HUB_TOKEN}" ]]; then
+	huggingface-cli login --token "${HF_TOKEN:-$HUGGING_FACE_HUB_TOKEN}" --add-to-git-credential
+else
+	echo "HF_TOKEN or HUGGING_FACE_HUB_TOKEN not set; skipping login"
+fi
+
+# Login to WanDB
+if [[ -n "${WANDB_API_KEY:-$WANDB_TOKEN}" ]]; then
+	wandb login "${WANDB_API_KEY:-$WANDB_TOKEN}"
+else
+	echo "WANDB_API_KEY or WANDB_TOKEN not set; skipping login"
+fi
 
 # 🫡
 sleep infinity

documentation/DOCKER.md

@@ -39,7 +39,7 @@ This command sets up the container with GPU access and maps the SSH port for ext
 To facilitate integration with external tools, the container supports environment variables for Huggingface and WandB tokens. Pass these at runtime as follows:
 
 ```bash
-docker run --gpus all -e HUGGING_FACE_HUB_TOKEN='your_token' -e WANDB_TOKEN='your_token' -it -p 22:22 simpletuner
+docker run --gpus all -e HF_TOKEN='your_token' -e WANDB_API_KEY='your_token' -it -p 22:22 simpletuner
 ```
 
 ### 4. Data Volumes
@@ -98,8 +98,8 @@ services:
       - "[path to your datasets]:/datasets"
       - "[path to your configs]:/workspace/SimpleTuner/config"
     environment:
-      HUGGING_FACE_HUB_TOKEN: [your hugging face token]
-      WANDB_TOKEN: [your wanddb token]
+      HF_TOKEN: [your hugging face token]
+      WANDB_API_KEY: [your wanddb token]
     command: ["tail", "-f", "/dev/null"]
     deploy:
       resources:
@@ -155,4 +155,4 @@ services:
 ### General Advice
 
 - **Logs and Output**: Review the container logs and output for any error messages or warnings that can provide more context on the issue.
-- **Documentation and Forums**: Consult the Docker and NVIDIA CUDA documentation for more detailed troubleshooting advice. Community forums and issue trackers related to the specific software or dependencies you are using can also be valuable resources.
+- **Documentation and Forums**: Consult the Docker and NVIDIA CUDA documentation for more detailed troubleshooting advice. Community forums and issue trackers related to the specific software or dependencies you are using can also be valuable resources.