attempt at launching small fast in CI, add tqdm_loggable (#719)

.github/workflows/docker-base-image.yaml

@@ -1,9 +1,12 @@
 name: Build and Push Docker TPU Images
 
 on:
- push:
- branches:
- - main
+ workflow_run:
+ workflows: ["Run Tests"]
+ types:
+ - completed
+ branches: [main]
+ workflow_dispatch:
 
 jobs:
  build:

.github/workflows/launch_small_fast.yaml

@@ -0,0 +1,72 @@
+name: Launch Llama 2 Small Fast
+
+on:
+ workflow_run:
+ workflows: ["Build and Push Docker TPU Images"]
+ types:
+ - completed
+ branches: [main, "experiment/*"]
+# pull_request:
+ workflow_dispatch:
+
+jobs:
+ test:
+ if: (github.event.pull_request.head.repo.full_name == github.repository)
+ runs-on: ubuntu-latest
+ env:
+ TPU_ZONE: "us-central2-b"
+ TPU_TYPE: "v4-32"
+
+ steps:
+ - name: Checkout code
+ uses: actions/checkout@v2
+
+ - name: Set up Google Cloud SDK
+ uses: google-github-actions/setup-gcloud@v1
+ with:
+ project_id: ${{ secrets.GCP_PROJECT_ID }}
+
+ - name: Authenticate to Google Cloud
+ uses: google-github-actions/auth@v1
+ with:
+ credentials_json: ${{ secrets.GCP_SA_KEY }}
+
+ - name: Configure Google Cloud
+ run: |
+ gcloud config set project ${{ secrets.GCP_PROJECT_ID }}
+ REGION=${TPU_ZONE%-*}
+ echo "$REGION"
+ gcloud auth configure-docker $REGION-docker.pkg.dev
+
+ - name: Install locally
+ run: |
+ python -m pip install --upgrade pip
+ pip install -e .[test] "jax[cpu]==0.4.30"
+
+ - name: Launch Small Fast TPU Train LM job
+ run: |
+ export TPU_NAME=small-fast-${{ github.run_id }}
+ export WANDB_API_KEY=${{ secrets.WANDB_API_KEY }}
+ export RUN_ID=small_fast_${{ github.run_id }}
+ export HF_TOKEN=${{ secrets.HF_TOKEN }}
+ 
+ cat > .config <<EOF
+ env:
+ WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
+ WANDB_ENTITY: stanford-mercury
+ WANDB_PROJECT: levanter
+ HF_TOKEN: ${{ secrets.HF_TOKEN }}
+ EOF
+ 
+ python infra/launch.py -e CI 1 --foreground --tpu_name ${TPU_NAME} --run_id $RUN_ID --zone ${TPU_ZONE} --tpu_type ${TPU_TYPE} --preemptible -- \
+ python -m levanter.main.train_lm \
+ --config_path config/llama_small_fast.yaml \
+ --trainer.checkpointer.base_path gs://levanter-checkpoints/llama-itest/ \
+ --trainer.checkpointer.save_interval 10m
+ --trainer.num_train_steps 10000
+
+ - name: Cleanup
+ if: ${{ always() }}
+ run: |
+ export TPU_NAME=small-fast-${{ github.run_id }}
+ gcloud compute tpus queued-resources delete $TPU_NAME --zone ${TPU_ZONE} --quiet --force

infra/launch.py

@@ -110,6 +110,8 @@ def main():
  else:
  raise ValueError(f"Unknown docker registry: {registry}")
 
+ failure = None
+
  for i in range(retries + 1):
  try:
  launch_job(
@@ -128,24 +130,39 @@ def main():
  print(f"Error running command {e.cmd}")
  if i < retries - 1:
  print("Retrying... %d/%d" % (i + 1, retries))
+ else:
+ print("Retries exhausted. Raising error.")
+ print(f"Error running command {e.cmd}")
+ print(f"Output: {e.output}")
+ failure = e
  else:
  print("Job finished with no error.")
  break
 
- if autodelete:
- print("Autodelete is set to True. Tearing down machine...")
- levanter.infra.tpus.run_command(
- "gcloud",
- "alpha",
- "compute",
- "tpus",
- "queued-resources",
- "delete",
- tpu_name,
- "--quiet",
- f"--zone={zone}",
- "--force",
- )
+ try:
+ if autodelete:
+ print("Autodelete is set to True. Tearing down machine...")
+ levanter.infra.tpus.run_command(
+ "gcloud",
+ "alpha",
+ "compute",
+ "tpus",
+ "queued-resources",
+ "delete",
+ tpu_name,
+ "--quiet",
+ f"--zone={zone}",
+ "--force",
+ )
+ except Exception as e:
+ print(f"Error tearing down TPU {tpu_name}: {e}")
+ if failure:
+ raise failure
+ else:
+ raise e
+
+ if failure:
+ raise failure
 
 
 if __name__ == "__main__":

pyproject.toml

@@ -50,6 +50,7 @@ dependencies = [
  "filelock~=3.13",
  # "ai2-olmo",
  "async-lru~=2.0",
+ "tqdm-loggable>=0.2"
 ]
 
 [project.urls]

src/levanter/callbacks.py

@@ -8,11 +8,13 @@
 import threading
 import time
 import warnings
+from datetime import timedelta
 from typing import Callable, Optional
 
 import humanfriendly
 import jax
-from tqdm import tqdm
+from tqdm_loggable import tqdm_logging
+from tqdm_loggable.auto import tqdm
 
 import levanter.tracker
 from levanter.data import DataLoader
@@ -39,7 +41,9 @@ def eval_loss_loop(loss_fn, model, dataset, max_batches: Optional[int] = None, n
  else:
  desc = "eval"
 
+ _tqdm_logging_one_time_setup()
  pbar = tqdm(dataset, desc=desc, position=1, leave=False, total=max_batches)
+
  iter_ = iter(pbar)
  while True:
  time_in = time.time()
@@ -186,6 +190,8 @@ def pbar_logger(iterable=None, desc="train", **tqdm_mkwargs):
  kwargs["desc"] = desc
  if "iterable" not in kwargs:
  kwargs["iterable"] = iterable
+
+ _tqdm_logging_one_time_setup()
  pbar = tqdm(**kwargs)
 
  def update_pbar(step: StepInfo):
@@ -359,3 +365,14 @@ def compute_and_viz_log_probs(step: StepInfo):
  wandb.log({"log_probs": wandb.Html(path)}, step=step.step)
 
  return compute_and_viz_log_probs
+
+
+_did_tqdm_logging_one_time_setup = False
+
+
+def _tqdm_logging_one_time_setup():
+ global _did_tqdm_logging_one_time_setup
+ if _did_tqdm_logging_one_time_setup:
+ return
+ _did_tqdm_logging_one_time_setup = True
+ tqdm_logging.tqdm_logging.set_log_rate(timedelta(seconds=60))

src/levanter/eval.py

@@ -8,8 +8,8 @@
 import jax.numpy as jnp
 import jmp
 import numpy as np
-import tqdm
 from jax.sharding import Mesh
+from tqdm_loggable.auto import tqdm
 
 import haliax as hax
 from haliax.partitioning import ResourceMapping
@@ -300,7 +300,7 @@ def evaluate(self, m: LmHeadModel):
  iterator = LoadingTimeTrackerIterator(self.loader)
  n = 0
 
- for batch, tags in tqdm.tqdm(iterator, "eval"):
+ for batch, tags in tqdm(iterator, "eval"):
  state = self.accum_for_batch(m, state, batch, tags)
  n += 1
 

src/levanter/infra/docker.py

@@ -65,7 +65,12 @@ def read(fd):
 
  return b"".join(output)
  else:
- return subprocess.check_output(argv, stderr=subprocess.STDOUT)
+ try:
+ return subprocess.check_output(argv, stderr=subprocess.STDOUT)
+ except subprocess.CalledProcessError as e:
+ # print the output if the command failed, reraising the exception
+ print(e.output.decode())
+ raise e
 
 
 def configure_gcp_docker(project_id, region, repository):

src/levanter/infra/tpus.py

@@ -49,6 +49,7 @@ def list_tpus(zone):
  "list",
  f"--zone={zone}",
  "--format=json(name.basename(), state)",
+ "--quiet",
  ]
  )
  )
@@ -68,6 +69,7 @@ def describe_tpu(tpu_name, zone):
  tpu_name,
  f"--zone={zone}",
  "--format=json(name.basename(), state)",
+ "--quiet",
  ],
  stderr=subprocess.DEVNULL,
  )
@@ -77,6 +79,8 @@ def describe_tpu(tpu_name, zone):
 
 
 def start_tpu_vm(tpu_name, *, tpu_type, capacity_type, version, zone, node_count):
+ # ensure alpha is enabled
+ run_command("gcloud", "components", "install", "alpha", "--quiet")
  if version is None:
  version = "tpu-ubuntu2204-base"
  tpu_stat = describe_tpu(tpu_name, zone)
@@ -196,17 +200,31 @@ def run_command(*args, **kwargs):
 
 def add_ssh_key(ssh_key_filename):
  # format 3072 SHA256:... key-name (RSA)
- key_hash = subprocess.check_output(["ssh-keygen", "-lf", ssh_key_filename]).decode("utf-8").split()[1]
- existing_keys = subprocess.check_output(["ssh-add", "-l"]).decode("utf-8").split("\n")
- for key in existing_keys:
- if key_hash in key:
- return
+ try:
+ key_hash = (
+ subprocess.check_output(["ssh-keygen", "-lf", ssh_key_filename], stderr=subprocess.STDOUT)
+ .decode("utf-8")
+ .split()[1]
+ )
+ existing_keys = (
+ subprocess.check_output(["ssh-add", "-l"], stderr=subprocess.STDOUT).decode("utf-8").split("\n")
+ )
+ for key in existing_keys:
+ if key_hash in key:
+ return
 
- subprocess.check_call(["ssh-add", ssh_key_filename])
+ subprocess.check_call(["ssh-add", ssh_key_filename])
+ except subprocess.CalledProcessError:
+ raise
 
 
 def tpu_ssh(tpu_name, zone, node_count, *args, ignore_failure=False):
- add_ssh_key(os.path.expanduser("~/.ssh/google_compute_engine"))
+ try:
+ add_ssh_key(os.path.expanduser("~/.ssh/google_compute_engine"))
+ except subprocess.CalledProcessError as e:
+ print("Failed to add ssh key. This may lead to problems.", e)
+ pass
+
  try:
  if node_count > 1:
  return _tpu_ssh_multislice(tpu_name, zone, node_count, *args, ignore_failure=ignore_failure)
@@ -219,6 +237,7 @@ def tpu_ssh(tpu_name, zone, node_count, *args, ignore_failure=False):
  "tpu-vm",
  "ssh",
  tpu_name,
+ "--quiet",
  "--worker=all",
  f"--zone={zone}",
  "--command=%s" % " ".join(args),
@@ -243,6 +262,7 @@ def _tpu_ssh_multislice(tpu_name, zone, node_count, *args, ignore_failure=False)
  "ssh",
  f"{tpu_name}-{i}",
  "--worker=all",
+ "--quiet",
  f"--zone={zone}",
  "--command=%s" % " ".join(args),
  )