Repository: pyro-ppl/pyro

Issue:
Allow testing of examples not parameterized by 'num_epochs'
The AIR example isn't run by the inference testing machinery, as it doesn't have a `--num-epochs` argument.

I could add such an argument, but this isn't a very natural parameterization for this example. Also, running an entire epoch for each example seems likely to be overkill in some cases, including AIR. (Assuming an epoch is a pass through the whole dataset.)

One idea is that the test should look for examples with a dedicated`--ci-test` argument or similar. This is a bit more flexible, though we might pay for this with a little extra code in each example.


Discussion:
@fritzo: Oh, wait, we'll need to hold off on this until the process of creating the data set used by the AIR example is automated.
Yes I saw that, I can add that in the test PR, unless you think that's tricky.
@fritzo: I'm not sure what you have in mind, but if you can make it work, then that will be great!

Resolving patch:
diff --git a/examples/air/main.py b/examples/air/main.py
--- a/examples/air/main.py
+++ b/examples/air/main.py
@@ -8,9 +8,11 @@
 """
 
 import os
+import sys
 import time
 import argparse
 from functools import partial
+from subprocess import check_call
 import numpy as np
 
 import torch
@@ -96,9 +98,10 @@
 # Load data.
 infile = './data/multi_mnist_train_uint8.npz'
 if not os.path.exists(infile):
-    print('Could not find the dataset at {}'.format(infile))
-    print('Run "python multi_mnist.py" to generate it.')
-    exit()
+    print('Running multi_mnist.py to generate dataset at {}...'.format(infile))
+    multi_mnist_py = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'multi_mnist.py')
+    check_call([sys.executable, multi_mnist_py])
+    print('Finished running multi_mnist.py.')
 X_np = np.load(infile)['x'].astype(np.float32)
 X_np /= 255.0
 X = Variable(torch.from_numpy(X_np))


Test patch:
diff --git a/tests/test_examples.py b/tests/test_examples.py
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -1,12 +1,16 @@
 import os
 import sys
+from collections import OrderedDict
 from subprocess import check_call
 
 from tests.common import EXAMPLES_DIR, requires_cuda
 import pytest
 
+CPU_EXAMPLES = OrderedDict()
+CUDA_EXAMPLES = OrderedDict()
 
-def find_examples_matching(*substrings):
+
+def discover_examples():
     for root, dirs, files in os.walk(EXAMPLES_DIR):
         for basename in files:
             if not basename.endswith('.py'):
@@ -14,20 +18,32 @@ def find_examples_matching(*substrings):
             path = os.path.join(root, basename)
             with open(path) as f:
                 text = f.read()
-            if all(s in text for s in substrings):
-                yield os.path.relpath(path, EXAMPLES_DIR)
+            if '--num-epochs' in text:
+                args = ['--num-epochs=1']
+            elif '--num-steps' in text:
+                args = ['--num-steps=1']
+            else:
+                # Either this is not a main file, or we don't know how to run it cheaply.
+                continue
+            example = os.path.relpath(path, EXAMPLES_DIR)
+            CPU_EXAMPLES[example] = args
+            if '--cuda' in text:
+                CUDA_EXAMPLES[example] = [args] + ['--cuda']
+
+
+discover_examples()
 
 
 @pytest.mark.stage("test_examples")
-@pytest.mark.parametrize('example', find_examples_matching('--num-epochs'))
-def test_cpu(example):
+@pytest.mark.parametrize('example,args', CPU_EXAMPLES.items(), ids=list(CPU_EXAMPLES))
+def test_cpu(example, args):
     example = os.path.join(EXAMPLES_DIR, example)
-    check_call([sys.executable, example, '--num-epochs', '1'])
+    check_call([sys.executable, example] + args)
 
 
 @requires_cuda
 @pytest.mark.stage("test_examples")
-@pytest.mark.parametrize('example', find_examples_matching('--num-epochs', '--cuda'))
-def test_cuda(example):
+@pytest.mark.parametrize('example,args', CUDA_EXAMPLES.items(), ids=list(CUDA_EXAMPLES))
+def test_cuda(example, args):
     example = os.path.join(EXAMPLES_DIR, example)
-    check_call([sys.executable, example, '--num-epochs', '1', '--cuda'])
+    check_call([sys.executable, example] + args)