[BE] lint/test/docs on PR and push to main (#742)

* lint/test on PR and push to main * Run docs CI on every push to main * Fail docs when a notebook doesn't build * skip eval_metrics fix in tutorial_utils * when writing extxyz in docs, specific force as a column explicitly * task.dataset -> dataset.format * small fixes in config in tutorial * make train dataset optional if only used test dataset * more small fixes if you don't specify a train dataset * loss_fns -> loss_functions in tutorial * eval_metrics -> evaluation_metrics in OCP tutorial * FixAtoms as int * don't clear train dataset if present but src is missing * invert train dataset load logic * fix logic * bool mask * require docs test on main * pull_request_review to trigger doc build * use FixAtoms fix from stress-relaxations branch * sneak in a push trigger to test docs * upload only html artifact * sneaky push * fix artifact path * fix job names * no deploy docs on push * add devcontainer settings * small fix to legacy tutorial * small fix to legacy tutorial * another small tutorial fix * add rest of tutorial back * clean up logic in load_datasets * typo * logic typo in load_datasets * convert all md to ipynb on launch * add minimum spec on devcontainer to make sure everything runs * include subdirectories in conversion * 758 fix * minor fixes * relax otf * typo * try removing comments in build_docs yml --------- Co-authored-by: zulissimeta <[email protected]> Co-authored-by: Zack Ulissi <[email protected]> Co-authored-by: Muhammed Shuaibi <[email protected]>
FAIR-Chem · Jul 13, 2024 · aacb883 · aacb883
1 parent 1442177
commit aacb883
Show file tree

Hide file tree

Showing 17 changed files with 180 additions and 88 deletions.
diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
@@ -0,0 +1,22 @@
+{
+    "image": "mcr.microsoft.com/devcontainers/anaconda:0.205.0-3",
+    "postCreateCommand": "bash ./.devcontainer/postCreateCommand.sh",
+    "forwardPorts": [
+        8888
+    ],
+    "postStartCommand": "nohup bash -c 'jupyter notebook --ip=0.0.0.0 --port=8888 --allow-root --NotebookApp.token='' &'",
+    // Configure tool-specific properties.
+    "customizations": {
+        // Configure properties specific to VS Code.
+        "vscode": {
+            // Add the IDs of extensions you want installed when the container is created.
+            "extensions": [
+                "ms-python.python",
+                "ms-python.vscode-pylance"
+            ]
+        }
+    },
+    "hostRequirements": {
+        "memory": "16gb",
+    }
+}
diff --git a/.devcontainer/postCreateCommand.sh b/.devcontainer/postCreateCommand.sh
@@ -0,0 +1,11 @@
+#!/bin/bash
+if [ -f packages/requirements.txt ]; then pip install -r packages/requirements.txt; fi
+if [ -f packages/requirements-optional.txt ]; then pip install -r packages/requirements-optional.txt; fi
+pip install -e packages/fairchem-core[dev]
+pip install -e packages/fairchem-data-oc[dev]
+pip install -e packages/fairchem-demo-ocpapi[dev]
+pip install -e packages/fairchem-applications-cattsunami
+pip install jupytext
+
+# Convert all .md docs to ipynb for easy viewing in vscode later!
+find ./docs -name '*.md' -exec jupytext --to ipynb {} \;
diff --git a/.github/workflows/docs.yml → .github/workflows/build_docs.yml b/.github/workflows/docs.yml → .github/workflows/build_docs.yml
@@ -1,12 +1,14 @@
-name: documentation
+# build the documentation and upload the built artifact
+name: build documentation
 
 on:
   workflow_call:
   workflow_dispatch:
+  pull_request_review:
+    types: [submitted, edited]
 
-# This job installs dependencies, build the book, and pushes it to `gh-pages`
 jobs:
-  build-and-deploy:
+  build:
     runs-on: ubuntu-latest
 
     steps:
@@ -36,17 +38,8 @@ jobs:
       run: |
         jupyter-book build docs
 
-    # Deploy the book's HTML to gh-pages branch # TODO remove once ODAC link updated
-    - name: Deploy to ghpages branch
-      uses: peaceiris/actions-gh-pages@v3
+    - name: Upload documentation artifact
+      uses: actions/upload-artifact@v4
       with:
-        github_token: ${{ secrets.GITHUB_TOKEN }}
-        publish_dir: docs/_build/html
-
-    - name: Deploy to fair-chem.github.io
-      uses: peaceiris/actions-gh-pages@v4
-      with:
-        deploy_key: ${{ secrets.ACTIONS_DEPLOY_KEY }}
-        external_repository: FAIR-Chem/fair-chem.github.io
-        publish_branch: gh-pages
-        publish_dir: docs/_build/html
+        name: docs-html
+        path: docs/_build/html/*
diff --git a/.github/workflows/deploy_docs.yml b/.github/workflows/deploy_docs.yml
@@ -0,0 +1,34 @@
+name: build and deploy docs
+
+on:
+  workflow_call:
+  workflow_dispatch:
+
+jobs:
+  build:
+    uses: ./.github/workflows/build_docs.yml
+
+  deploy:
+    needs: build
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Download docs artifact
+        uses: actions/download-artifact@v4
+        with:
+          name: docs-html
+          path: docs-html/
+
+      - name: Deploy to ghpages branch
+        uses: peaceiris/actions-gh-pages@v3
+        with:
+          github_token: ${{ secrets.GITHUB_TOKEN }}
+          publish_dir: docs-html
+
+      - name: Deploy to fair-chem.github.io
+        uses: peaceiris/actions-gh-pages@v4
+        with:
+          deploy_key: ${{ secrets.ACTIONS_DEPLOY_KEY }}
+          external_repository: FAIR-Chem/fair-chem.github.io
+          publish_branch: gh-pages
+          publish_dir: docs-html
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
@@ -1,8 +1,11 @@
 name: lint
 
 on:
-  push:
+  workflow_call:
   workflow_dispatch:
+  pull_request:
+    branches: [main]
+  push:
 
 jobs:
   lint:

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -2,9 +2,11 @@ name: tests
 
 on:
   workflow_call:
+  workflow_dispatch:
   pull_request:
     branches: [main]
-  workflow_dispatch:
+  push:
+    branches: [main]
 
 jobs:
   test:

diff --git a/docs/_config.yml b/docs/_config.yml
@@ -52,3 +52,5 @@ sphinx:
     autosummary_generate: True
     autoapi_dirs: ['../src/fairchem/core','../src/fairchem/data','../src/fairchem/applications/AdsorbML/adsorbml','../src/fairchem/applications/cattsunami','../src/fairchem/demo/ocpapi']
     autoapi_python_use_implicit_namespaces: True
+    nb_execution_show_tb: True
+    nb_execution_raise_on_error: True
diff --git a/docs/core/fine-tuning/fine-tuning-oxides.md b/docs/core/fine-tuning/fine-tuning-oxides.md
@@ -207,21 +207,23 @@ yml = generate_yml_config(checkpoint_path, 'config.yml',
                            'optim.loss_force', # the checkpoint setting causes an error
                            'dataset', 'test_dataset', 'val_dataset'],
                    update={'gpus': 1,
-                           'task.dataset': 'ase_db',
                            'optim.eval_every': 10,
                            'optim.max_epochs': 1,
                            'optim.batch_size': 4,
                             'logger':'tensorboard', # don't use wandb!
                            # Train data
                            'dataset.train.src': 'train.db',
+                           'dataset.train.format': 'ase_db',
                            'dataset.train.a2g_args.r_energy': True,
                            'dataset.train.a2g_args.r_forces': True,
                             # Test data - prediction only so no regression
                            'dataset.test.src': 'test.db',
+                           'dataset.test.format': 'ase_db',
                            'dataset.test.a2g_args.r_energy': False,
                            'dataset.test.a2g_args.r_forces': False,
                            # val data
                            'dataset.val.src': 'val.db',
+                           'dataset.val.format': 'ase_db',
                            'dataset.val.a2g_args.r_energy': True,
                            'dataset.val.a2g_args.r_forces': True,
                           })

diff --git a/docs/core/inference.md b/docs/core/inference.md
@@ -47,6 +47,12 @@ with ase.db.connect('full_data.db') as full_db:
       if 'tag' in atoms.info['key_value_pairs']:
         atoms.info['key_value_pairs']['tag'] = int(atoms.info['key_value_pairs']['tag'])
 
+      for key in atoms.info["key_value_pairs"]:
+        if atoms.info["key_value_pairs"][key] == "True":
+            atoms.info["key_value_pairs"][key] = True
+        elif atoms.info["key_value_pairs"][key] == "False":
+            atoms.info["key_value_pairs"][key] = False
+
       subset_db.write(atoms, **atoms.info['key_value_pairs'])
 ```
 
@@ -78,11 +84,11 @@ yml = generate_yml_config(checkpoint_path, 'config.yml',
                            'dataset', 'slurm'],
                    update={'amp': True,
                            'gpus': 1,
-                           'task.dataset': 'ase_db',
                            'task.prediction_dtype': 'float32',
                            'logger':'tensorboard', # don't use wandb!
                             # Test data - prediction only so no regression
                            'dataset.test.src': 'data.db',
+                           'dataset.test.format': 'ase_db',
                            'dataset.test.a2g_args.r_energy': False,
                            'dataset.test.a2g_args.r_forces': False,
                            'dataset.test.select_args.selection': 'natoms>5,xc=PBE',

diff --git a/docs/legacy_tutorials/OCP_Tutorial.md b/docs/legacy_tutorials/OCP_Tutorial.md
@@ -302,7 +302,7 @@ dyn.run(fmax=0, steps=100)
 traj = ase.io.read("data/toy_c3h8_relax.traj", ":")
 
 # convert traj format to extxyz format (used by OC20 dataset)
-columns = (['symbols','positions', 'move_mask', 'tags'])
+columns = (['symbols','positions', 'move_mask', 'tags', 'forces'])
 with open('data/toy_c3h8_relax.extxyz','w') as f:
     extxyz.write_xyz(f, traj, columns=columns)
 ```
@@ -925,8 +925,8 @@ trainer = OCPTrainer(
     dataset=dataset,
     optimizer=optimizer,
     outputs={},
-    loss_fns={},
-    eval_metrics={},
+    loss_functions={},
+    evaluation_metrics={},
     name="s2ef",
     identifier="S2EF-example",
     run_dir=".", # directory to save results if is_debug=False. Prediction files are saved here so be careful not to override!
@@ -1007,8 +1007,8 @@ pretrained_trainer = OCPTrainer(
     dataset=dataset,
     optimizer=optimizer,
     outputs={},
-    loss_fns={},
-    eval_metrics={},
+    loss_functions={},
+    evaluation_metrics={},
     name="s2ef",
     identifier="S2EF-val-example",
     run_dir="./", # directory to save results if is_debug=False. Prediction files are saved here so be careful not to override!
@@ -1211,8 +1211,8 @@ energy_trainer = OCPTrainer(
     dataset=dataset,
     optimizer=optimizer,
     outputs={},
-    loss_fns={},
-    eval_metrics={},
+    loss_functions={},
+    evaluation_metrics={},
     name="is2re",
     identifier="IS2RE-example",
     run_dir="./", # directory to save results if is_debug=False. Prediction files are saved here so be careful not to override!
@@ -1289,8 +1289,8 @@ pretrained_energy_trainer = OCPTrainer(
     dataset=dataset,
     optimizer=optimizer,
     outputs={},
-    loss_fns={},
-    eval_metrics={},
+    loss_functions={},
+    evaluation_metrics={},
     name="is2re",
     identifier="IS2RE-val-example",
     run_dir="./", # directory to save results if is_debug=False. Prediction files are saved here so be careful not to override!
@@ -1427,7 +1427,6 @@ task = {
       'grad_input': 'atomic forces',
       'train_on_free_atoms': True,
       'eval_on_free_atoms': True,
-      'relax_dataset': {"src": relax_dataset},
       'write_pos': True,
       'relaxation_steps': 200,
       'num_relaxation_batches': 1,
@@ -1489,10 +1488,11 @@ optimizer = {
     'force_coefficient': 100,
 }
 # Dataset
-dataset = [
-  {'src': train_src, 'normalize_labels': False}, # train set
-  {'src': val_src}, # val set (optional)
-]
+dataset = {
+  'train': {'src': train_src, 'normalize_labels': False}, # train set
+  'val': {'src': val_src}, # val set (optional)
+  'relax': {"src": relax_dataset},
+}
 ```
 
 
@@ -1514,8 +1514,8 @@ trainer = OCPTrainer(
     dataset=dataset,
     optimizer=optimizer,
     outputs={},
-    loss_fns={},
-    eval_metrics={},
+    loss_functions={},
+    evaluation_metrics={},
     name="s2ef",
     identifier="is2rs-example",
     run_dir="./", # directory to save results if is_debug=False. Prediction files are saved here so be careful not to override!
@@ -1904,8 +1904,8 @@ trainer = OCPTrainer(
     dataset=dataset,
     optimizer=optimizer,
     outputs={},
-    loss_fns={},
-    eval_metrics={},
+    loss_functions={},
+    evaluation_metrics={},
     name="s2ef",
     identifier="S2EF-simple",
     run_dir="./", # directory to save results if is_debug=False. Prediction files are saved here so be careful not to override!
@@ -1980,8 +1980,8 @@ trainer = OCPTrainer(
     dataset=dataset,
     optimizer=optimizer,
     outputs={},
-    loss_fns={},
-    eval_metrics={},
+    loss_functions={},
+    evaluation_metrics={},
     name="s2ef",
     identifier="S2EF-gemnet-t",
     run_dir="./", # directory to save results if is_debug=False. Prediction files are saved here so be careful not to override!

diff --git a/docs/legacy_tutorials/data_visualization.md b/docs/legacy_tutorials/data_visualization.md
@@ -101,7 +101,7 @@ dyn.run(fmax=0, steps=100)
 traj = ase.io.read("data/toy_c3h8_relax.traj", ":")
 
 # convert traj format to extxyz format (used by OC20 dataset)
-columns = (['symbols','positions', 'move_mask', 'tags'])
+columns = (['symbols','positions', 'move_mask', 'tags', 'forces'])
 with open('data/toy_c3h8_relax.extxyz','w') as f:
     extxyz.write_xyz(f, traj, columns=columns)
     

diff --git a/docs/tutorials/advanced/fine-tuning-in-python.md b/docs/tutorials/advanced/fine-tuning-in-python.md
@@ -79,21 +79,23 @@ yml = generate_yml_config(checkpoint_path, 'config.yml',
                            'optim.loss_force', # the checkpoint setting causes an error
                            'dataset', 'test_dataset', 'val_dataset'],
                    update={'gpus': 1,
-                           'task.dataset': 'ase_db',
                            'optim.eval_every': 10,
                            'optim.max_epochs': 1,
                            'optim.batch_size': 4,
                            'logger': 'tensorboard', # don't use wandb unless you already are logged in
                            # Train data
                            'dataset.train.src': 'train.db',
+                           'dataset.train.format': 'ase_db',
                            'dataset.train.a2g_args.r_energy': True,
                            'dataset.train.a2g_args.r_forces': True,
                             # Test data - prediction only so no regression
                            'dataset.test.src': 'test.db',
+                           'dataset.test.format': 'ase_db',
                            'dataset.test.a2g_args.r_energy': False,
                            'dataset.test.a2g_args.r_forces': False,
                            # val data
                            'dataset.val.src': 'val.db',
+                           'dataset.val.format': 'ase_db',
                            'dataset.val.a2g_args.r_energy': True,
                            'dataset.val.a2g_args.r_forces': True,
                           })

diff --git a/src/fairchem/core/common/relaxation/ase_utils.py b/src/fairchem/core/common/relaxation/ase_utils.py
@@ -38,7 +38,7 @@ def batch_to_atoms(batch):
     n_systems = batch.natoms.shape[0]
     natoms = batch.natoms.tolist()
     numbers = torch.split(batch.atomic_numbers, natoms)
-    fixed = torch.split(batch.fixed, natoms)
+    fixed = torch.split(batch.fixed.to(torch.bool), natoms)
     forces = torch.split(batch.force, natoms)
     positions = torch.split(batch.pos, natoms)
     tags = torch.split(batch.tags, natoms)

diff --git a/src/fairchem/core/common/relaxation/optimizers/lbfgs_torch.py b/src/fairchem/core/common/relaxation/optimizers/lbfgs_torch.py
@@ -142,7 +142,7 @@ def run(self, fmax, steps):
                 traj_fl = Path(self.traj_dir / f"{name}.traj_tmp", mode="w")
                 traj_fl.rename(traj_fl.with_suffix(".traj"))
 
-        self.batch.y, self.batch.force = self.get_energy_and_forces(
+        self.batch.energy, self.batch.force = self.get_energy_and_forces(
             apply_constraint=False
         )
         return self.batch
@@ -211,7 +211,7 @@ def determine_step(dr):
         self.f0 = forces
 
     def write(self, energy, forces, update_mask) -> None:
-        self.batch.y, self.batch.force = energy, forces
+        self.batch.energy, self.batch.force = energy, forces
         atoms_objects = batch_to_atoms(self.batch)
         update_mask_ = torch.split(update_mask, self.batch.natoms.tolist())
         for atm, traj, mask in zip(atoms_objects, self.trajectories, update_mask_):

diff --git a/src/fairchem/core/common/tutorial_utils.py b/src/fairchem/core/common/tutorial_utils.py
@@ -174,10 +174,6 @@ def nested_set(dic, keys, value):
         keys = _key.split(".")
         nested_set(config, keys, update[_key])
 
-    # TODO : Do not rename keys in utils.py when reading a config
-    config["evaluation_metrics"] = config["eval_metrics"]
-    config.pop("eval_metrics")
-
     out = dump(config)
     with open(yml, "wb") as f:
         f.write(out.encode("utf-8"))