From cd8a8c19b765601858a9521d8e8b84f07c081259 Mon Sep 17 00:00:00 2001
From: Romil Bhardwaj <romil.bhardwaj@berkeley.edu>
Date: Wed, 20 Mar 2024 17:56:18 -0700
Subject: [PATCH] [examples] Update examples using deprecated `SKYPILOT_JOB_ID
 ` (#3344)

Update examples
---
 examples/tensorflow_distributed/tf_distributed.yaml | 4 ++--
 llm/vicuna/train.yaml                               | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/examples/tensorflow_distributed/tf_distributed.yaml b/examples/tensorflow_distributed/tf_distributed.yaml
index ac76cec98ce..beb6ad4b96e 100644
--- a/examples/tensorflow_distributed/tf_distributed.yaml
+++ b/examples/tensorflow_distributed/tf_distributed.yaml
@@ -53,12 +53,12 @@ run: |
               'index': node_rank
           }
       }
-  with open(f'/tmp/{os.environ.get("SKYPILOT_JOB_ID")}', 'w') as f:
+  with open(f'/tmp/{os.environ.get("SKYPILOT_TASK_ID")}', 'w') as f:
       json.dump(tf_config, f)
   EOF
   
   # Read and set TF_CONFIG from file
-  export TF_CONFIG=$(cat /tmp/$SKYPILOT_JOB_ID)
+  export TF_CONFIG=$(cat /tmp/$SKYPILOT_TASK_ID)
   echo $TF_CONFIG
   
   # ======== Run the training script ========
diff --git a/llm/vicuna/train.yaml b/llm/vicuna/train.yaml
index 4afd93f1478..c577561e858 100644
--- a/llm/vicuna/train.yaml
+++ b/llm/vicuna/train.yaml
@@ -100,7 +100,7 @@ run: |
     --fsdp_transformer_layer_cls_to_wrap 'LlamaDecoderLayer' \
     --tf32 True \
     --model_max_length ${SEQ_LEN} \
-    --run_name $SKYPILOT_JOB_ID \
+    --run_name $SKYPILOT_TASK_ID \
     --gradient_checkpointing True \
     --lazy_preprocess True