Add AZURE_OPENAI_API_KEY to auto env vars (#514)

Signed-off-by: Igor Gitman <igitman@nvidia.com>

Author: Kipok

docs/basics/cluster-configs.md

@@ -47,6 +47,8 @@ export WANDB_API_KEY=...
 export HF_TOKEN=...
 # only needed if running inference with OpenAI models
 export OPENAI_API_KEY=...
+# only needed if running inference with Azure OpenAI models
+export AZURE_OPENAI_API_KEY=...
 # only needed if running inference with Nvidia NIM models
 export NVIDIA_API_KEY=...
 ```

nemo_skills/pipeline/setup.py

@@ -96,8 +96,11 @@ def setup():
 
         env_vars = typer.prompt(
             "\nYou can also specify any environment variables that you want to be accessible in containers.\n"
-            "Can either define just the name (we take value from the current environment), or name=value to use a fixed value.\n"
-            "By default we will always pass WANDB_API_KEY, NVIDIA_API_KEY, OPENAI_API_KEY, HF_TOKEN, so you don't need to list those.\n"
+            "Can either define just the name (we take value from the current environment), "
+            "or name=value to use a fixed value.\n"
+            "By default we will always pass "
+            "WANDB_API_KEY, NVIDIA_API_KEY, AZURE_OPENAI_API_KEY, OPENAI_API_KEY, HF_TOKEN, "
+            "so you don't need to list those.\n"
             "What other environment variables would you like to define? (comma separated)",
             default="",
         )

nemo_skills/pipeline/utils/cluster.py

@@ -63,7 +63,7 @@ def get_env_variables(cluster_config):
     - `required_env_vars` - list of required environment variables
     - `env_vars` - list of optional environment variables
 
-    WANDB_API_KEY, NVIDIA_API_KEY, OPENAI_API_KEY, and HF_TOKEN are always added if they exist.
+    WANDB_API_KEY, NVIDIA_API_KEY, AZURE_OPENAI_API_KEY, OPENAI_API_KEY, and HF_TOKEN are always added if they exist.
 
     Args:
         cluster_config: cluster config dictionary
@@ -101,7 +101,13 @@ def get_env_variables(cluster_config):
     # It is fine to have these as always optional even if they are required for some configs
     # Assume it is required, then this will override the value set above with the same
     # value, assuming it has not been updated externally between these two calls
-    always_optional_env_vars = ["WANDB_API_KEY", "NVIDIA_API_KEY", "OPENAI_API_KEY", "HF_TOKEN"]
+    always_optional_env_vars = [
+        "WANDB_API_KEY",
+        "NVIDIA_API_KEY",
+        "AZURE_OPENAI_API_KEY",
+        "OPENAI_API_KEY",
+        "HF_TOKEN",
+    ]
     default_factories = {
         "HF_TOKEN": lambda: str(get_token()),
     }