From 80e8629621fb172d2dca2beffe7fc21b17af8318 Mon Sep 17 00:00:00 2001
From: Aryan Gupta <guptaaryan16@gmail.com>
Date: Mon, 21 Aug 2023 23:06:40 +0530
Subject: [PATCH 1/3] Updated UI for ArgParser options in TrainingTab.vue

---
 src/components/TabTraining.vue         | 21 +++++++++++++++++++++
 src/metadata/metadata.json             |  7 +++++++
 src/store.js                           |  1 +
 src/templates/template-common/utils.py | 10 ++++++++++
 4 files changed, 39 insertions(+)

diff --git a/src/components/TabTraining.vue b/src/components/TabTraining.vue
index bdfcba41..237e066c 100644
--- a/src/components/TabTraining.vue
+++ b/src/components/TabTraining.vue
@@ -1,6 +1,25 @@
 <template>
   <div class="tab training">
     <h1>Training Options</h1>
+
+    We are using yaml files to set up the training configuration (see
+    config.yaml). Additional arguments given to the main script can be parsed
+    using one of the tools provided below:
+
+    <h2>Argument Parser</h2>
+
+    <ul>
+      <li>
+        <a href="https://docs.python.org/3/library/argparse.html">Argparse</a> -
+        is a python built-in tool to handle command-line arguments
+      </li>
+    </ul>
+    <FormSelect
+      :label="argparser.description"
+      :options="argparser.options"
+      :saveKey="argparser.name"
+      :defaultV="argparser.default"
+    />
     <h2 class="training">Deterministic Training</h2>
     <FormCheckbox
       :label="deterministic.description"
@@ -64,6 +83,7 @@ export default {
   components: { FormCheckbox, FormInput, FormRadio, FormSelect },
   setup() {
     const {
+      argparser,
       deterministic,
       backend,
       torchrun,
@@ -81,6 +101,7 @@ export default {
     })
     return {
       store,
+      argparser,
       deterministic,
       backend,
       torchrun,
diff --git a/src/metadata/metadata.json b/src/metadata/metadata.json
index 5c07af09..855ed330 100644
--- a/src/metadata/metadata.json
+++ b/src/metadata/metadata.json
@@ -1,5 +1,12 @@
 {
   "training": {
+    "argparser": {
+      "name": "argparser",
+      "type": "array",
+      "description": "Select the argument parser for training",
+      "options": ["argparse"],
+      "default": "argparse"
+    },
     "deterministic": {
       "name": "deterministic",
       "type": "checkbox",
diff --git a/src/store.js b/src/store.js
index 6ebb3dfe..6b95114e 100644
--- a/src/store.js
+++ b/src/store.js
@@ -38,6 +38,7 @@ export const msg = reactive({
 export const default_config = {
   template: '',
   include_test: false,
+  argparser: 'argparse',
   output_dir: './logs',
   log_every_iters: 10,
   logger: 'tensorboard',
diff --git a/src/templates/template-common/utils.py b/src/templates/template-common/utils.py
index 48375097..2084b16e 100644
--- a/src/templates/template-common/utils.py
+++ b/src/templates/template-common/utils.py
@@ -1,5 +1,9 @@
 import logging
+
+#::: if ((it.argparser == 'argparse')) { :::#
 from argparse import ArgumentParser
+
+#::: } :::#
 from datetime import datetime
 from logging import Logger
 from pathlib import Path
@@ -37,6 +41,9 @@
 from omegaconf import DictConfig, OmegaConf
 
 
+#::: if ((it.argparser == 'argparse')) { :::#
+
+
 def get_default_parser():
     parser = ArgumentParser()
     parser.add_argument("config", type=Path, help="Config file path")
@@ -64,6 +71,9 @@ def setup_config(parser=None):
     return DictConfig(config)
 
 
+#::: } :::#
+
+
 def log_metrics(engine: Engine, tag: str) -> None:
     """Log `engine.state.metrics` with given `engine` and `tag`.
 

From 8ef4bf87237d8f8e9e1991e72609d5884c768ccd Mon Sep 17 00:00:00 2001
From: Aryan Gupta <guptaaryan16@gmail.com>
Date: Tue, 22 Aug 2023 13:27:50 +0530
Subject: [PATCH 2/3] Introduce CI changes for argparse UI option

---
 .github/workflows/ci.yml                |   9 +-
 __tests__/text-classification.spec.js   | 344 ++++++++++++------------
 __tests__/vision-classification.spec.js | 341 +++++++++++------------
 __tests__/vision-dcgan.spec.js          | 336 ++++++++++++-----------
 __tests__/vision-segmentation.spec.js   | 343 +++++++++++------------
 package.json                            |   2 +-
 scripts/run_code_style.sh               |   6 +-
 scripts/run_tests.sh                    |   8 +-
 8 files changed, 720 insertions(+), 669 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index bef566f1..ecbc1ec6 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -80,9 +80,14 @@ jobs:
           pip install tqdm
 
       - run: pnpm build
-      - run: pnpm test:ci
+
+      - name: 'pnpm test:ci ${{ matrix.template }}'
+        env:
+          NPM_CONFIG: ${{ matrix.template }}
+        run: pnpm test:ci
+
       - run: sh ./scripts/run_tests.sh unzip
-      - run: pnpm dist_lint
+      - run: pnpm dist_lint ${{ matrix.template }}
 
       - name: 'Run ${{ matrix.template }} ${{ matrix.test }}'
         run: sh ./scripts/run_tests.sh ${{ matrix.test }} ${{ matrix.template }}
diff --git a/__tests__/text-classification.spec.js b/__tests__/text-classification.spec.js
index bf72b277..41bf4e47 100644
--- a/__tests__/text-classification.spec.js
+++ b/__tests__/text-classification.spec.js
@@ -24,169 +24,181 @@ afterEach(async () => {
   await context.close()
 })
 
-test('text classification simple', async () => {
-  await page.selectOption('select', 'template-text-classification')
-
-  await page.waitForSelector('text=README.md')
-
-  await page.click('text=Loggers')
-  await page.click('text=config.yaml')
-
-  // TODO: simplify the downloadPromise calls
-  // Here we are trying to wait for 2 seconds before clicking on the `Code` and `Download Zip` button
-  const downloadPromise = await page
-    .waitForEvent('download', { timeout: 2000 })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      // these catch calls are required to make sure if CI fails initially then we can have something to rely for further tests
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-
-  await downloadPromise.saveAs('./dist-tests/text-classification-simple.zip')
-})
-
-test('text classification all', async () => {
-  await page.selectOption('select', 'template-text-classification')
-
-  await page.check('#include_test-checkbox')
-  expect(await page.isChecked('#include_test-checkbox')).toBeTruthy()
-
-  await page.waitForSelector('text=README.md')
-  await page.click('text=Training')
-
-  await page.check('#deterministic-checkbox')
-  expect(await page.isChecked('#deterministic-checkbox')).toBeTruthy()
-
-  await page.click('text=Handlers')
-
-  await page.check('#save_training-checkbox')
-  expect(await page.isChecked('#save_training-checkbox')).toBeTruthy()
-
-  await page.check('#save_evaluation-checkbox')
-  expect(await page.isChecked('#save_evaluation-checkbox')).toBeTruthy()
-
-  await page.fill('#filename_prefix-input-text', 'training')
-  expect(await page.$eval('#filename_prefix-input-text', (e) => e.value)).toBe(
-    'training'
-  )
-
-  await page.fill('#save_every_iters-input-number', '2')
-  expect(
-    await page.$eval('#save_every_iters-input-number', (e) => e.value)
-  ).toBe('2')
-
-  await page.fill('#n_saved-input-number', '2')
-  expect(await page.$eval('#n_saved-input-number', (e) => e.value)).toBe('2')
-
-  await page.check('#terminate_on_nan-checkbox')
-  expect(await page.isChecked('#terminate_on_nan-checkbox')).toBeTruthy()
-
-  await page.fill('#patience-input-number', '2')
-  expect(await page.$eval('#patience-input-number', (e) => e.value)).toBe('2')
-
-  await page.fill('#limit_sec-input-number', '60')
-  expect(await page.$eval('#limit_sec-input-number', (e) => e.value)).toBe('60')
-
-  await page.click('text=Loggers')
-  await page.click('text=config.yaml')
-
-  const downloadPromise = await page
-    .waitForEvent('download', { timeout: 2000 })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-  await downloadPromise.saveAs('./dist-tests/text-classification-all.zip')
-})
-
-test('text classification launch', async () => {
-  await page.selectOption('select', 'template-text-classification')
-
-  await page.waitForSelector('text=README.md')
-  await page.click('text=Training')
-
-  await page.check('#use_dist-checkbox')
-  expect(await page.isChecked('#use_dist-checkbox')).toBeTruthy()
-
-  expect(await page.isChecked('#dist-torchrun-radio')).toBeTruthy()
-
-  await page.click('text=Loggers')
-  await page.click('text=config.yaml')
-
-  const downloadPromise = await page
-    .waitForEvent('download', { timeout: 2000 })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-
-  await downloadPromise.saveAs('./dist-tests/text-classification-launch.zip')
-})
-
-test('text classification spawn', async () => {
-  await page.selectOption('select', 'template-text-classification')
-
-  await page.waitForSelector('text=README.md')
-  await page.click('text=Training')
-
-  await page.check('#use_dist-checkbox')
-  expect(await page.isChecked('#use_dist-checkbox')).toBeTruthy()
-
-  await page.check('#dist-spawn-radio')
-  expect(await page.isChecked('#dist-spawn-radio')).toBeTruthy()
-
-  await page.click('text=Loggers')
-  await page.click('text=config.yaml')
-
-  const downloadPromise = await page
-    .waitForEvent('download', { timeout: 2000 })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-  await downloadPromise.saveAs('./dist-tests/text-classification-spawn.zip')
-})
+const parser = ['argparse']
+for (const name of parser) {
+  test(`text-classification simple ${name}`, async () => {
+    await page.selectOption('select', 'template-text-classification')
+
+    await page.waitForSelector('text=README.md')
+
+    await page.click('text=Loggers')
+    await page.click('text=config.yaml')
+    await page.getByText('Training', { exact: true }).click()
+
+    // TODO: simplify the downloadPromise calls
+    // Here we are trying to wait for 2 seconds before clicking on the `Code` and `Download Zip` button
+    const downloadPromise = await page
+      .waitForEvent('download', { timeout: 2000 })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        // these catch calls are required to make sure if CI fails initially then we can have something to rely for further tests
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+
+    await downloadPromise.saveAs(
+      `./dist-tests/text-classification-simple-${name}.zip`
+    )
+  })
+
+  test(`text-classification all ${name}`, async () => {
+    await page.selectOption('select', 'template-text-classification')
+
+    await page.check('#include_test-checkbox')
+    expect(await page.isChecked('#include_test-checkbox')).toBeTruthy()
+
+    await page.waitForSelector('text=README.md')
+    await page.getByText('Training', { exact: true }).click()
+    await page.check('#deterministic-checkbox')
+    expect(await page.isChecked('#deterministic-checkbox')).toBeTruthy()
+
+    await page.click('text=Handlers')
+
+    await page.check('#save_training-checkbox')
+    expect(await page.isChecked('#save_training-checkbox')).toBeTruthy()
+
+    await page.check('#save_evaluation-checkbox')
+    expect(await page.isChecked('#save_evaluation-checkbox')).toBeTruthy()
+
+    await page.fill('#filename_prefix-input-text', 'training')
+    expect(
+      await page.$eval('#filename_prefix-input-text', (e) => e.value)
+    ).toBe('training')
+
+    await page.fill('#save_every_iters-input-number', '2')
+    expect(
+      await page.$eval('#save_every_iters-input-number', (e) => e.value)
+    ).toBe('2')
+
+    await page.fill('#n_saved-input-number', '2')
+    expect(await page.$eval('#n_saved-input-number', (e) => e.value)).toBe('2')
+
+    await page.check('#terminate_on_nan-checkbox')
+    expect(await page.isChecked('#terminate_on_nan-checkbox')).toBeTruthy()
+
+    await page.fill('#patience-input-number', '2')
+    expect(await page.$eval('#patience-input-number', (e) => e.value)).toBe('2')
+
+    await page.fill('#limit_sec-input-number', '60')
+    expect(await page.$eval('#limit_sec-input-number', (e) => e.value)).toBe(
+      '60'
+    )
+
+    await page.click('text=Loggers')
+    await page.click('text=config.yaml')
+
+    const downloadPromise = await page
+      .waitForEvent('download', { timeout: 2000 })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+    await downloadPromise.saveAs(
+      `./dist-tests/text-classification-all-${name}.zip`
+    )
+  })
+
+  test(`text-classification launch ${name}`, async () => {
+    await page.selectOption('select', 'template-text-classification')
+
+    await page.waitForSelector('text=README.md')
+    await page.getByText('Training', { exact: true }).click()
+
+    await page.check('#use_dist-checkbox')
+    expect(await page.isChecked('#use_dist-checkbox')).toBeTruthy()
+
+    expect(await page.isChecked('#dist-torchrun-radio')).toBeTruthy()
+
+    await page.click('text=Loggers')
+    await page.click('text=config.yaml')
+
+    const downloadPromise = await page
+      .waitForEvent('download', { timeout: 2000 })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+
+    await downloadPromise.saveAs(
+      `./dist-tests/text-classification-launch-${name}.zip`
+    )
+  })
+
+  test('text-classification spawn', async () => {
+    await page.selectOption('select', 'template-text-classification')
+
+    await page.waitForSelector('text=README.md')
+    await page.getByText('Training', { exact: true }).click()
+    await page.check('#use_dist-checkbox')
+    expect(await page.isChecked('#use_dist-checkbox')).toBeTruthy()
+
+    await page.check('#dist-spawn-radio')
+    expect(await page.isChecked('#dist-spawn-radio')).toBeTruthy()
+
+    await page.click('text=Loggers')
+    await page.click('text=config.yaml')
+
+    const downloadPromise = await page
+      .waitForEvent('download', { timeout: 2000 })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+    await downloadPromise.saveAs(
+      `./dist-tests/text-classification-spawn-${name}.zip`
+    )
+  })
+}
diff --git a/__tests__/vision-classification.spec.js b/__tests__/vision-classification.spec.js
index 4502fbce..9f2850fd 100644
--- a/__tests__/vision-classification.spec.js
+++ b/__tests__/vision-classification.spec.js
@@ -24,168 +24,179 @@ afterEach(async () => {
   await context.close()
 })
 
-test('vision classification simple', async () => {
-  await page.selectOption('select', 'template-vision-classification')
-
-  await page.waitForSelector('text=README.md')
-
-  await page.click('text=Loggers')
-  await page.click('text=config.yaml')
-
-  // TODO: simplify the downloadPromise calls
-  // Here we are trying to wait for 2 seconds before clicking on the `Code` and `Download Zip` button
-  const downloadPromise = await page
-    .waitForEvent('download', { timeout: 2000 })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      // these catch calls are required to make sure if CI fails initially then we can have something to rely for further tests
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-
-  await downloadPromise.saveAs('./dist-tests/vision-classification-simple.zip')
-})
-
-test('vision classification all', async () => {
-  await page.selectOption('select', 'template-vision-classification')
-
-  await page.check('#include_test-checkbox')
-  expect(await page.isChecked('#include_test-checkbox')).toBeTruthy()
-
-  await page.waitForSelector('text=README.md')
-  await page.click('text=Training')
-
-  await page.check('#deterministic-checkbox')
-  expect(await page.isChecked('#deterministic-checkbox')).toBeTruthy()
-
-  await page.click('text=Handlers')
-
-  await page.check('#save_training-checkbox')
-  expect(await page.isChecked('#save_training-checkbox')).toBeTruthy()
-
-  await page.check('#save_evaluation-checkbox')
-  expect(await page.isChecked('#save_evaluation-checkbox')).toBeTruthy()
-
-  await page.fill('#filename_prefix-input-text', 'training')
-  expect(await page.$eval('#filename_prefix-input-text', (e) => e.value)).toBe(
-    'training'
-  )
-
-  await page.fill('#save_every_iters-input-number', '2')
-  expect(
-    await page.$eval('#save_every_iters-input-number', (e) => e.value)
-  ).toBe('2')
-
-  await page.fill('#n_saved-input-number', '2')
-  expect(await page.$eval('#n_saved-input-number', (e) => e.value)).toBe('2')
-
-  await page.check('#terminate_on_nan-checkbox')
-  expect(await page.isChecked('#terminate_on_nan-checkbox')).toBeTruthy()
-
-  await page.fill('#patience-input-number', '2')
-  expect(await page.$eval('#patience-input-number', (e) => e.value)).toBe('2')
-
-  await page.fill('#limit_sec-input-number', '60')
-  expect(await page.$eval('#limit_sec-input-number', (e) => e.value)).toBe('60')
-
-  await page.click('text=Loggers')
-  await page.click('text=config.yaml')
-
-  const downloadPromise = await page
-    .waitForEvent('download', { timeout: 2000 })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-  await downloadPromise.saveAs('./dist-tests/vision-classification-all.zip')
-})
-
-test('vision classification launch', async () => {
-  await page.selectOption('select', 'template-vision-classification')
-
-  await page.waitForSelector('text=README.md')
-  await page.click('text=Training')
-
-  await page.check('#use_dist-checkbox')
-  expect(await page.isChecked('#use_dist-checkbox')).toBeTruthy()
-
-  expect(await page.isChecked('#dist-torchrun-radio')).toBeTruthy()
-
-  await page.click('text=Loggers')
-  await page.click('text=config.yaml')
-
-  const downloadPromise = await page
-    .waitForEvent('download', { timeout: 2000 })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-  await downloadPromise.saveAs('./dist-tests/vision-classification-launch.zip')
-})
-
-test('vision classification spawn', async () => {
-  await page.selectOption('select', 'template-vision-classification')
-
-  await page.waitForSelector('text=README.md')
-  await page.click('text=Training')
-
-  await page.check('#use_dist-checkbox')
-  expect(await page.isChecked('#use_dist-checkbox')).toBeTruthy()
-
-  await page.check('#dist-spawn-radio')
-  expect(await page.isChecked('#dist-spawn-radio')).toBeTruthy()
-
-  await page.click('text=Loggers')
-  await page.click('text=config.yaml')
-
-  const downloadPromise = await page
-    .waitForEvent('download', { timeout: 2000 })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-  await downloadPromise.saveAs('./dist-tests/vision-classification-spawn.zip')
-})
+const parser = ['argparse']
+for (const name of parser) {
+  test(`vision-classification simple ${name}`, async () => {
+    await page.selectOption('select', 'template-vision-classification')
+
+    await page.waitForSelector('text=README.md')
+
+    await page.click('text=Loggers')
+    await page.click('text=config.yaml')
+    await page.getByText('Training', { exact: true }).click()
+
+    // TODO: simplify the downloadPromise calls
+    // Here we are trying to wait for 2 seconds before clicking on the `Code` and `Download Zip` button
+    const downloadPromise = await page
+      .waitForEvent('download', { timeout: 2000 })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        // these catch calls are required to make sure if CI fails initially then we can have something to rely for further tests
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+
+    await downloadPromise.saveAs(
+      `./dist-tests/vision-classification-simple-${name}.zip`
+    )
+  })
+
+  test(`vision-classification all ${name}`, async () => {
+    await page.selectOption('select', 'template-vision-classification')
+
+    await page.check('#include_test-checkbox')
+    expect(await page.isChecked('#include_test-checkbox')).toBeTruthy()
+
+    await page.waitForSelector('text=README.md')
+    await page.getByText('Training', { exact: true }).click()
+    await page.check('#deterministic-checkbox')
+    expect(await page.isChecked('#deterministic-checkbox')).toBeTruthy()
+
+    await page.click('text=Handlers')
+
+    await page.check('#save_training-checkbox')
+    expect(await page.isChecked('#save_training-checkbox')).toBeTruthy()
+
+    await page.check('#save_evaluation-checkbox')
+    expect(await page.isChecked('#save_evaluation-checkbox')).toBeTruthy()
+
+    await page.fill('#filename_prefix-input-text', 'training')
+    expect(
+      await page.$eval('#filename_prefix-input-text', (e) => e.value)
+    ).toBe('training')
+
+    await page.fill('#save_every_iters-input-number', '2')
+    expect(
+      await page.$eval('#save_every_iters-input-number', (e) => e.value)
+    ).toBe('2')
+
+    await page.fill('#n_saved-input-number', '2')
+    expect(await page.$eval('#n_saved-input-number', (e) => e.value)).toBe('2')
+
+    await page.check('#terminate_on_nan-checkbox')
+    expect(await page.isChecked('#terminate_on_nan-checkbox')).toBeTruthy()
+
+    await page.fill('#patience-input-number', '2')
+    expect(await page.$eval('#patience-input-number', (e) => e.value)).toBe('2')
+
+    await page.fill('#limit_sec-input-number', '60')
+    expect(await page.$eval('#limit_sec-input-number', (e) => e.value)).toBe(
+      '60'
+    )
+
+    await page.click('text=Loggers')
+    await page.click('text=config.yaml')
+
+    const downloadPromise = await page
+      .waitForEvent('download', { timeout: 2000 })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+    await downloadPromise.saveAs(
+      `./dist-tests/vision-classification-all-${name}.zip`
+    )
+  })
+
+  test('vision-classification launch ${name}', async () => {
+    await page.selectOption('select', 'template-vision-classification')
+
+    await page.waitForSelector('text=README.md')
+    await page.getByText('Training', { exact: true }).click()
+    await page.check('#use_dist-checkbox')
+    expect(await page.isChecked('#use_dist-checkbox')).toBeTruthy()
+
+    expect(await page.isChecked('#dist-torchrun-radio')).toBeTruthy()
+
+    await page.click('text=Loggers')
+    await page.click('text=config.yaml')
+
+    const downloadPromise = await page
+      .waitForEvent('download', { timeout: 2000 })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+    await downloadPromise.saveAs(
+      `./dist-tests/vision-classification-launch-${name}.zip`
+    )
+  })
+
+  test(`vision-classification spawn ${name}`, async () => {
+    await page.selectOption('select', 'template-vision-classification')
+
+    await page.waitForSelector('text=README.md')
+    await page.getByText('Training', { exact: true }).click()
+    await page.check('#use_dist-checkbox')
+    expect(await page.isChecked('#use_dist-checkbox')).toBeTruthy()
+
+    await page.check('#dist-spawn-radio')
+    expect(await page.isChecked('#dist-spawn-radio')).toBeTruthy()
+
+    await page.click('text=Loggers')
+    await page.click('text=config.yaml')
+
+    const downloadPromise = await page
+      .waitForEvent('download', { timeout: 2000 })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+    await downloadPromise.saveAs(
+      `./dist-tests/vision-classification-spawn-${name}.zip`
+    )
+  })
+}
diff --git a/__tests__/vision-dcgan.spec.js b/__tests__/vision-dcgan.spec.js
index 5cbe8a03..25edb21e 100644
--- a/__tests__/vision-dcgan.spec.js
+++ b/__tests__/vision-dcgan.spec.js
@@ -24,167 +24,175 @@ afterEach(async () => {
   await context.close()
 })
 
-test('vision dcgan simple', async () => {
-  await page.selectOption('select', 'template-vision-dcgan')
-
-  await page.waitForSelector('text=README.md')
-
-  await page.click('text=Loggers')
-  await page.click('text=config.yaml')
-
-  // TODO: simplify the downloadPromise calls
-  // Here we are trying to wait for 2 seconds before clicking on the `Code` and `Download Zip` button
-  const downloadPromise = await page
-    .waitForEvent('download', { timeout: 2000 })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      // these catch calls are required to make sure if CI fails initially then we can have something to rely for further tests
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-  await downloadPromise.saveAs('./dist-tests/vision-dcgan-simple.zip')
-})
-
-test('vision dcgan all', async () => {
-  await page.selectOption('select', 'template-vision-dcgan')
-
-  await page.check('#include_test-checkbox')
-  expect(await page.isChecked('#include_test-checkbox')).toBeTruthy()
-
-  await page.waitForSelector('text=README.md')
-  await page.click('text=Training')
-
-  await page.check('#deterministic-checkbox')
-  expect(await page.isChecked('#deterministic-checkbox')).toBeTruthy()
-
-  await page.click('text=Handlers')
-
-  await page.check('#save_training-checkbox')
-  expect(await page.isChecked('#save_training-checkbox')).toBeTruthy()
-
-  await page.check('#save_evaluation-checkbox')
-  expect(await page.isChecked('#save_evaluation-checkbox')).toBeTruthy()
-
-  await page.fill('#filename_prefix-input-text', 'training')
-  expect(await page.$eval('#filename_prefix-input-text', (e) => e.value)).toBe(
-    'training'
-  )
-
-  await page.fill('#save_every_iters-input-number', '2')
-  expect(
-    await page.$eval('#save_every_iters-input-number', (e) => e.value)
-  ).toBe('2')
-
-  await page.fill('#n_saved-input-number', '2')
-  expect(await page.$eval('#n_saved-input-number', (e) => e.value)).toBe('2')
-
-  await page.check('#terminate_on_nan-checkbox')
-  expect(await page.isChecked('#terminate_on_nan-checkbox')).toBeTruthy()
-
-  await page.fill('#patience-input-number', '2')
-  expect(await page.$eval('#patience-input-number', (e) => e.value)).toBe('2')
-
-  await page.fill('#limit_sec-input-number', '60')
-  expect(await page.$eval('#limit_sec-input-number', (e) => e.value)).toBe('60')
-
-  await page.click('text=Loggers')
-  await page.click('text=config.yaml')
-
-  const downloadPromise = await page
-    .waitForEvent('download', { timeout: 2000 })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-  await downloadPromise.saveAs('./dist-tests/vision-dcgan-all.zip')
-})
-
-test('vision dcgan launch', async () => {
-  await page.selectOption('select', 'template-vision-dcgan')
-
-  await page.waitForSelector('text=README.md')
-  await page.click('text=Training')
-
-  await page.check('#use_dist-checkbox')
-  expect(await page.isChecked('#use_dist-checkbox')).toBeTruthy()
-
-  expect(await page.isChecked('#dist-torchrun-radio')).toBeTruthy()
-
-  await page.click('text=Loggers')
-  await page.click('text=config.yaml')
-
-  const downloadPromise = await page
-    .waitForEvent('download', { timeout: 2000 })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-  await downloadPromise.saveAs('./dist-tests/vision-dcgan-launch.zip')
-})
-
-test('vision dcgan spawn', async () => {
-  await page.selectOption('select', 'template-vision-dcgan')
-
-  await page.waitForSelector('text=README.md')
-  await page.click('text=Training')
-
-  await page.check('#use_dist-checkbox')
-  expect(await page.isChecked('#use_dist-checkbox')).toBeTruthy()
-
-  await page.check('#dist-spawn-radio')
-  expect(await page.isChecked('#dist-spawn-radio')).toBeTruthy()
-
-  await page.click('text=Loggers')
-  await page.click('text=config.yaml')
-
-  const downloadPromise = await page
-    .waitForEvent('download', { timeout: 2000 })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-  await downloadPromise.saveAs('./dist-tests/vision-dcgan-spawn.zip')
-})
+const parser = ['argparse']
+for (const name of parser) {
+  test(`vision-dcgan simple ${name}`, async () => {
+    await page.selectOption('select', 'template-vision-dcgan')
+
+    await page.waitForSelector('text=README.md')
+
+    await page.click('text=Loggers')
+    await page.click('text=config.yaml')
+    await page.getByText('Training', { exact: true }).click()
+
+    // TODO: simplify the downloadPromise calls
+    // Here we are trying to wait for 2 seconds before clicking on the `Code` and `Download Zip` button
+    const downloadPromise = await page
+      .waitForEvent('download', { timeout: 2000 })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        // these catch calls are required to make sure if CI fails initially then we can have something to rely for further tests
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+    await downloadPromise.saveAs(`./dist-tests/vision-dcgan-simple-${name}.zip`)
+  })
+
+  test(`vision-dcgan all ${name}`, async () => {
+    await page.selectOption('select', 'template-vision-dcgan')
+
+    await page.check('#include_test-checkbox')
+    expect(await page.isChecked('#include_test-checkbox')).toBeTruthy()
+
+    await page.waitForSelector('text=README.md')
+    await page.getByText('Training', { exact: true }).click()
+    await page.check('#deterministic-checkbox')
+    expect(await page.isChecked('#deterministic-checkbox')).toBeTruthy()
+
+    await page.click('text=Handlers')
+
+    await page.check('#save_training-checkbox')
+    expect(await page.isChecked('#save_training-checkbox')).toBeTruthy()
+
+    await page.check('#save_evaluation-checkbox')
+    expect(await page.isChecked('#save_evaluation-checkbox')).toBeTruthy()
+
+    await page.fill('#filename_prefix-input-text', 'training')
+    expect(
+      await page.$eval('#filename_prefix-input-text', (e) => e.value)
+    ).toBe('training')
+
+    await page.fill('#save_every_iters-input-number', '2')
+    expect(
+      await page.$eval('#save_every_iters-input-number', (e) => e.value)
+    ).toBe('2')
+
+    await page.fill('#n_saved-input-number', '2')
+    expect(await page.$eval('#n_saved-input-number', (e) => e.value)).toBe('2')
+
+    await page.check('#terminate_on_nan-checkbox')
+    expect(await page.isChecked('#terminate_on_nan-checkbox')).toBeTruthy()
+
+    await page.fill('#patience-input-number', '2')
+    expect(await page.$eval('#patience-input-number', (e) => e.value)).toBe('2')
+
+    await page.fill('#limit_sec-input-number', '60')
+    expect(await page.$eval('#limit_sec-input-number', (e) => e.value)).toBe(
+      '60'
+    )
+
+    await page.click('text=Loggers')
+    await page.click('text=config.yaml')
+
+    const downloadPromise = await page
+      .waitForEvent('download', { timeout: 2000 })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+    await downloadPromise.saveAs(`./dist-tests/vision-dcgan-all-${name}.zip`)
+  })
+
+  test(`vision-dcgan launch ${name}`, async () => {
+    await page.selectOption('select', 'template-vision-dcgan')
+
+    await page.waitForSelector('text=README.md')
+    await page.getByText('Training', { exact: true }).click()
+    await page.check('#use_dist-checkbox')
+    expect(await page.isChecked('#use_dist-checkbox')).toBeTruthy()
+
+    expect(await page.isChecked('#dist-torchrun-radio')).toBeTruthy()
+
+    await page.click('text=Loggers')
+    await page.click('text=config.yaml')
+
+    const downloadPromise = await page
+      .waitForEvent('download', { timeout: 2000 })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+    await downloadPromise.saveAs(`./dist-tests/vision-dcgan-launch-${name}.zip`)
+  })
+
+  test(`vision-dcgan spawn ${name}`, async () => {
+    await page.selectOption('select', 'template-vision-dcgan')
+
+    await page.waitForSelector('text=README.md')
+    await page.getByText('Training', { exact: true }).click()
+    await page
+      .getByRole('combobox', {
+        name: 'Select the argument parser for training'
+      })
+      .selectOption(`${name}`)
+    await page.check('#use_dist-checkbox')
+    expect(await page.isChecked('#use_dist-checkbox')).toBeTruthy()
+
+    await page.check('#dist-spawn-radio')
+    expect(await page.isChecked('#dist-spawn-radio')).toBeTruthy()
+
+    await page.click('text=Loggers')
+    await page.click('text=config.yaml')
+
+    const downloadPromise = await page
+      .waitForEvent('download', { timeout: 2000 })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+    await downloadPromise.saveAs(`./dist-tests/vision-dcgan-spawn-${name}.zip`)
+  })
+}
diff --git a/__tests__/vision-segmentation.spec.js b/__tests__/vision-segmentation.spec.js
index 60baa3a7..547e5f3a 100644
--- a/__tests__/vision-segmentation.spec.js
+++ b/__tests__/vision-segmentation.spec.js
@@ -24,169 +24,180 @@ afterEach(async () => {
   await context.close()
 })
 
-test('vision segmentation simple', async () => {
-  await page.selectOption('select', 'template-vision-segmentation')
-
-  await page.waitForSelector('text=README.md')
-
-  await page.click('text=Loggers')
-  await page.click('text=config.yaml')
-
-  // TODO: simplify the downloadPromise calls
-  // Here we are trying to wait for 2 seconds before clicking on the `Code` and `Download Zip` button
-  const downloadPromise = await page
-    .waitForEvent('download', { timeout: 2000 })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      // these catch calls are required to make sure if CI fails initially then we can have something to rely for further tests
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-  await downloadPromise.saveAs('./dist-tests/vision-segmentation-simple.zip')
-})
-
-test('vision segmentation all', async () => {
-  await page.selectOption('select', 'template-vision-segmentation')
-
-  await page.check('#include_test-checkbox')
-  expect(await page.isChecked('#include_test-checkbox')).toBeTruthy()
-
-  await page.waitForSelector('text=README.md')
-  await page.click('text=Training')
-
-  await page.check('#deterministic-checkbox')
-  expect(await page.isChecked('#deterministic-checkbox')).toBeTruthy()
-
-  await page.click('text=Handlers')
-
-  await page.check('#save_training-checkbox')
-  expect(await page.isChecked('#save_training-checkbox')).toBeTruthy()
-
-  await page.check('#save_evaluation-checkbox')
-  expect(await page.isChecked('#save_evaluation-checkbox')).toBeTruthy()
-
-  await page.fill('#filename_prefix-input-text', 'training')
-  expect(await page.$eval('#filename_prefix-input-text', (e) => e.value)).toBe(
-    'training'
-  )
-
-  await page.fill('#save_every_iters-input-number', '2')
-  expect(
-    await page.$eval('#save_every_iters-input-number', (e) => e.value)
-  ).toBe('2')
-
-  await page.fill('#n_saved-input-number', '2')
-  expect(await page.$eval('#n_saved-input-number', (e) => e.value)).toBe('2')
-
-  await page.check('#terminate_on_nan-checkbox')
-  expect(await page.isChecked('#terminate_on_nan-checkbox')).toBeTruthy()
-
-  await page.fill('#patience-input-number', '2')
-  expect(await page.$eval('#patience-input-number', (e) => e.value)).toBe('2')
-
-  await page.fill('#limit_sec-input-number', '60')
-  expect(await page.$eval('#limit_sec-input-number', (e) => e.value)).toBe('60')
-
-  await page.click('text=Loggers')
-  await page.click('text=config.yaml')
-
-  const downloadPromise = await page
-    .waitForEvent('download', { timeout: 2000 })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-
-  await downloadPromise.saveAs('./dist-tests/vision-segmentation-all.zip')
-})
-
-test('vision segmentation launch', async () => {
-  await page.selectOption('select', 'template-vision-segmentation')
-
-  await page.waitForSelector('text=README.md')
-  await page.click('text=Training')
-
-  await page.check('#use_dist-checkbox')
-  expect(await page.isChecked('#use_dist-checkbox')).toBeTruthy()
-
-  expect(await page.isChecked('#dist-torchrun-radio')).toBeTruthy()
-
-  await page.click('text=Loggers')
-  await page.click('text=config.yaml')
-
-  const downloadPromise = await page
-    .waitForEvent('download', { timeout: 2000 })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-  await downloadPromise.saveAs('./dist-tests/vision-segmentation-launch.zip')
-})
-
-test('vision segmentation spawn', async () => {
-  await page.selectOption('select', 'template-vision-segmentation')
-
-  await page.waitForSelector('text=README.md')
-  await page.click('text=Training')
-
-  await page.check('#use_dist-checkbox')
-  expect(await page.isChecked('#use_dist-checkbox')).toBeTruthy()
-
-  await page.check('#dist-spawn-radio')
-  expect(await page.isChecked('#dist-spawn-radio')).toBeTruthy()
-
-  await page.click('text=Loggers')
-  await page.click('text=config.yaml')
-
-  const downloadPromise = await page
-    .waitForEvent('download', { timeout: 2000 })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-    .catch(() => {
-      page.getByRole('button', { name: 'Code' }).click()
-      page.getByRole('button', { name: 'Download Zip' }).click()
-      return page.waitForEvent('download', { timeout: 2000 })
-    })
-
-  await downloadPromise.saveAs('./dist-tests/vision-segmentation-spawn.zip')
-})
+const parser = ['argparse', 'fire']
+for (const name of parser) {
+  test(`vision-segmentation simple ${name}`, async () => {
+    await page.selectOption('select', 'template-vision-segmentation')
+
+    await page.waitForSelector('text=README.md')
+
+    await page.click('text=Loggers')
+    await page.click('text=config.yaml')
+    await page.getByText('Training', { exact: true }).click()
+
+    // TODO: simplify the downloadPromise calls
+    // Here we are trying to wait for 2 seconds before clicking on the `Code` and `Download Zip` button
+    const downloadPromise = await page
+      .waitForEvent('download', { timeout: 2000 })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        // these catch calls are required to make sure if CI fails initially then we can have something to rely for further tests
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+    await downloadPromise.saveAs(
+      `./dist-tests/vision-segmentation-simple-${name}.zip`
+    )
+  })
+
+  test(`vision-segmentation all ${name}`, async () => {
+    await page.selectOption('select', 'template-vision-segmentation')
+
+    await page.check('#include_test-checkbox')
+    expect(await page.isChecked('#include_test-checkbox')).toBeTruthy()
+
+    await page.waitForSelector('text=README.md')
+    await page.getByText('Training', { exact: true }).click()
+    await page.check('#deterministic-checkbox')
+    expect(await page.isChecked('#deterministic-checkbox')).toBeTruthy()
+
+    await page.click('text=Handlers')
+
+    await page.check('#save_training-checkbox')
+    expect(await page.isChecked('#save_training-checkbox')).toBeTruthy()
+
+    await page.check('#save_evaluation-checkbox')
+    expect(await page.isChecked('#save_evaluation-checkbox')).toBeTruthy()
+
+    await page.fill('#filename_prefix-input-text', 'training')
+    expect(
+      await page.$eval('#filename_prefix-input-text', (e) => e.value)
+    ).toBe('training')
+
+    await page.fill('#save_every_iters-input-number', '2')
+    expect(
+      await page.$eval('#save_every_iters-input-number', (e) => e.value)
+    ).toBe('2')
+
+    await page.fill('#n_saved-input-number', '2')
+    expect(await page.$eval('#n_saved-input-number', (e) => e.value)).toBe('2')
+
+    await page.check('#terminate_on_nan-checkbox')
+    expect(await page.isChecked('#terminate_on_nan-checkbox')).toBeTruthy()
+
+    await page.fill('#patience-input-number', '2')
+    expect(await page.$eval('#patience-input-number', (e) => e.value)).toBe('2')
+
+    await page.fill('#limit_sec-input-number', '60')
+    expect(await page.$eval('#limit_sec-input-number', (e) => e.value)).toBe(
+      '60'
+    )
+
+    await page.click('text=Loggers')
+    await page.click('text=config.yaml')
+
+    const downloadPromise = await page
+      .waitForEvent('download', { timeout: 2000 })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+
+    await downloadPromise.saveAs(
+      `./dist-tests/vision-segmentation-all-${name}.zip`
+    )
+  })
+
+  test(`vision-segmentation launch ${name}`, async () => {
+    await page.selectOption('select', 'template-vision-segmentation')
+
+    await page.waitForSelector('text=README.md')
+    await page.getByText('Training', { exact: true }).click()
+    await page.check('#use_dist-checkbox')
+    expect(await page.isChecked('#use_dist-checkbox')).toBeTruthy()
+
+    expect(await page.isChecked('#dist-torchrun-radio')).toBeTruthy()
+
+    await page.click('text=Loggers')
+    await page.click('text=config.yaml')
+
+    const downloadPromise = await page
+      .waitForEvent('download', { timeout: 2000 })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+    await downloadPromise.saveAs(
+      `./dist-tests/vision-segmentation-launch-${name}.zip`
+    )
+  })
+
+  test(`vision-segmentation spawn ${name}`, async () => {
+    await page.selectOption('select', 'template-vision-segmentation')
+
+    await page.waitForSelector('text=README.md')
+    await page.getByText('Training', { exact: true }).click()
+    await page.check('#use_dist-checkbox')
+    expect(await page.isChecked('#use_dist-checkbox')).toBeTruthy()
+
+    await page.check('#dist-spawn-radio')
+    expect(await page.isChecked('#dist-spawn-radio')).toBeTruthy()
+
+    await page.click('text=Loggers')
+    await page.click('text=config.yaml')
+
+    const downloadPromise = await page
+      .waitForEvent('download', { timeout: 2000 })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+      .catch(() => {
+        page.getByRole('button', { name: 'Code' }).click()
+        page.getByRole('button', { name: 'Download Zip' }).click()
+        return page.waitForEvent('download', { timeout: 2000 })
+      })
+
+    await downloadPromise.saveAs(
+      `./dist-tests/vision-segmentation-spawn-${name}.zip`
+    )
+  })
+}
diff --git a/package.json b/package.json
index 04093da5..2b635f50 100644
--- a/package.json
+++ b/package.json
@@ -5,7 +5,7 @@
     "dev": "vite --port 5000",
     "build": "vite build",
     "serve": "vite preview",
-    "test": "rm -rf ./dist-tests && jest --color --runInBand",
+    "test": "rm -rf ./dist-tests && jest --color --runInBand -t ${NPM_CONFIG:=''}",
     "test:ci": "start-server-and-test --expect 200 serve http://127.0.0.1:5000 test",
     "release": "node scripts/release.js",
     "fmt": "prettier --write . && bash scripts/run_code_style.sh fmt",
diff --git a/scripts/run_code_style.sh b/scripts/run_code_style.sh
index f40be558..0d9d2fec 100755
--- a/scripts/run_code_style.sh
+++ b/scripts/run_code_style.sh
@@ -4,7 +4,11 @@ set -xeu
 
 if [ $1 == "dist_lint" ]; then
     # Check that ./dist-tests/ exists and code is unzipped
-    ls ./dist-tests/vision-classification-all/main.py
+    TEMP=${2:-vision-classification}
+
+    # for argparse
+    ls ./dist-tests/$TEMP-all-argparse/main.py
+    
     # Comment dist-tests in .gitignore to make black running on ./dist-tests folder
     sed -i "s/dist-tests/# dist-tests/g" .gitignore
 
diff --git a/scripts/run_tests.sh b/scripts/run_tests.sh
index 564a031c..1308d74b 100644
--- a/scripts/run_tests.sh
+++ b/scripts/run_tests.sh
@@ -15,7 +15,7 @@ unzip_all() {
 }
 
 run_simple() {
-  for dir in $(find ./dist-tests/$1-simple -type d)
+  for dir in $(find ./dist-tests/$1-simple-argparse -type d)
   do
     cd $dir
     python main.py ../../src/tests/ci-configs/$1-simple.yaml
@@ -24,7 +24,7 @@ run_simple() {
 }
 
 run_all() {
-  for dir in $(find ./dist-tests/$1-all -type d)
+  for dir in $(find ./dist-tests/$1-all-argparse -type d)
   do
     cd $dir
     pytest -vra --color=yes --tb=short test_*.py
@@ -34,7 +34,7 @@ run_all() {
 }
 
 run_launch() {
-  for dir in $(find ./dist-tests/$1-launch -type d)
+  for dir in $(find ./dist-tests/$1-launch-argparse -type d)
   do
     cd $dir
     torchrun --nproc_per_node 2 main.py ../../src/tests/ci-configs/$1-launch.yaml --backend gloo
@@ -43,7 +43,7 @@ run_launch() {
 }
 
 run_spawn() {
-  for dir in $(find ./dist-tests/$1-spawn -type d)
+  for dir in $(find ./dist-tests/$1-spawn-argparse -type d)
   do
     cd $dir
     python main.py ../../src/tests/ci-configs/$1-spawn.yaml --backend gloo

From 56e10add779ed3fb8bc91b7eacefb267304af802 Mon Sep 17 00:00:00 2001
From: Aryan Gupta <guptaaryan16@gmail.com>
Date: Tue, 22 Aug 2023 13:38:58 +0530
Subject: [PATCH 3/3] Remove bug in tests

---
 __tests__/vision-segmentation.spec.js | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/__tests__/vision-segmentation.spec.js b/__tests__/vision-segmentation.spec.js
index 547e5f3a..57343564 100644
--- a/__tests__/vision-segmentation.spec.js
+++ b/__tests__/vision-segmentation.spec.js
@@ -24,7 +24,7 @@ afterEach(async () => {
   await context.close()
 })
 
-const parser = ['argparse', 'fire']
+const parser = ['argparse']
 for (const name of parser) {
   test(`vision-segmentation simple ${name}`, async () => {
     await page.selectOption('select', 'template-vision-segmentation')