Fix notebooks (mosaicml#2446)

* fix autoresume with slashed directory

* Revert "fix autoresume with slashed directory"

This reverts commit 3dfb5f5.

revert

* fix

* remove print

examples/TPU_Training_in_composer.ipynb

@@ -120,7 +120,7 @@
    "source": [
     "from torchvision import datasets, transforms\n",
     "\n",
-    "data_directory = \"../data\"\n",
+    "data_directory = \"./data\"\n",
     "\n",
     "# Normalization constants\n",
     "mean = (0.507, 0.487, 0.441)\n",

examples/ffcv_dataloaders.ipynb

@@ -125,7 +125,7 @@
     "\n",
     "batch_size = 1024\n",
     "num_workers = 2\n",
-    "data_directory = \"/tmp\"\n",
+    "data_directory = \"./data\"\n",
     "\n",
     "cifar10_transforms = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean, std)])\n",
     "\n",

examples/getting_started.ipynb

@@ -125,7 +125,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "data_directory = \"../data\"\n",
+    "data_directory = \"./data\"\n",
     "\n",
     "# Normalization constants\n",
     "mean = (0.507, 0.487, 0.441)\n",

examples/training_with_submitit.ipynb

@@ -128,7 +128,7 @@
     "        [transforms.ToTensor(), transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010]),]\n",
     "    )\n",
     "    train_data = torchvision.datasets.CIFAR10(\n",
-    "        root=\".\", train=True, download=True, transform=cifar10_transform\n",
+    "        root=\"./data\", train=True, download=True, transform=cifar10_transform\n",
     "    )\n",
     "    train_loader = DataLoader(\n",
     "        dataset=train_data,\n",

tests/test_notebooks.py

@@ -4,12 +4,14 @@
 import glob
 import inspect
 import os
+from urllib.parse import urlparse
 
 import pytest
 import testbook
 from testbook.client import TestbookNotebookClient
 
 import composer
+from composer.utils.import_helpers import MissingConditionalImportError
 from tests.common import device
 
 nb_root = os.path.join(os.path.dirname(composer.__file__), '..', 'examples')
@@ -65,12 +67,14 @@ def modify_cell_source(tb: TestbookNotebookClient, notebook_name: str, cell_sour
     # This function is called before each cell is executed
     if notebook_name == 'functional_api':
         # avoid div by 0 errors with batch size of 1
-        cell_source = cell_source.replace('max_epochs = 5', 'max_epochs = 1')
+        cell_source = cell_source.replace('num_epochs = 5', 'num_epochs = 1')
         cell_source = cell_source.replace('acc_percent = 100 * num_right / eval_size', 'acc_percent = 1')
         cell_source = cell_source.replace('batch_size = 1024', 'batch_size = 64')
+        cell_source = cell_source.replace('download=True', 'download=False')
     if notebook_name == 'custom_speedup_methods':
         cell_source = cell_source.replace('resnet_56', 'resnet_9')
         cell_source = cell_source.replace('batch_size=1024', 'batch_size=64')
+        cell_source = cell_source.replace('download=True', 'download=False')
     if notebook_name == 'finetune_huggingface':
         cell_source = cell_source.replace(
             'sst2_dataset = datasets.load_dataset("glue", "sst2")',
@@ -81,10 +85,16 @@ def modify_cell_source(tb: TestbookNotebookClient, notebook_name: str, cell_sour
         cell_source = cell_source.replace('batch_size=32', 'batch_size=1')
     if notebook_name == 'early_stopping':
         cell_source = cell_source.replace('batch_size = 1024', 'batch_size = 64')
+        cell_source = cell_source.replace('download=True', 'download=False')
     if notebook_name == 'getting_started':
         cell_source = cell_source.replace('batch_size = 1024', 'batch_size = 64')
+        cell_source = cell_source.replace('download=True', 'download=False')
+    if notebook_name == 'auto_microbatching':
+        cell_source = cell_source.replace('download=True', 'download=False')
     if notebook_name == 'migrate_from_ptl':
         cell_source = cell_source.replace('batch_size=256', 'batch_size=64')
+        cell_source = cell_source.replace('download=True', 'download=False')
+
     return cell_source
 
 
@@ -94,6 +104,7 @@ def modify_cell_source(tb: TestbookNotebookClient, notebook_name: str, cell_sour
 def test_notebook(notebook: str, device: str, s3_bucket: str):
     trainer_monkeypatch_code = inspect.getsource(patch_notebooks)
     notebook_name = os.path.split(notebook)[-1][:-len('.ipynb')]
+
     if notebook_name == 'medical_image_segmentation':
         pytest.skip('Dataset is only available via kaggle; need to authenticate on ci/cd')
     if notebook_name == 'training_with_submitit':
@@ -113,6 +124,24 @@ def test_notebook(notebook: str, device: str, s3_bucket: str):
     if notebook_name == 'pretrain_finetune_huggingface':
         pytest.skip(
             "Error that is unreproducible locally: No module named 'transformers.models.mega.configuration_mega'")
+
+    try:
+        import boto3
+    except ImportError as e:
+        raise MissingConditionalImportError('streaming', 'boto3') from e
+
+    obj = urlparse('s3://mosaicml-internal-integration-testing/CIFAR-10/')
+    s3 = boto3.resource('s3')
+    bucket = s3.Bucket(obj.netloc)
+    files = bucket.objects.filter(Prefix=obj.path.lstrip('/'))
+    for file in files:
+        target = os.path.join(os.getcwd(), 'data', os.path.relpath(file.key, obj.path.lstrip('/')))
+        if not os.path.exists(target):
+            os.makedirs(os.path.dirname(target), exist_ok=True)
+        if file.key[-1] == '/':
+            continue
+        bucket.download_file(file.key, target)
+
     with testbook.testbook(notebook) as tb:
         tb.inject(trainer_monkeypatch_code)
         tb.inject('patch_notebooks()')