Fix nightly tests

.github/workflows/dev-ci.yaml

@@ -14,6 +14,10 @@ concurrency:
 defaults:
   run:
     shell: bash -l {0}
+
+env:
+  OE_LICENSE: ${{ github.workspace }}/oe_license.txt
+  DGL_HOME: ${{ github.workspace }}/dgl
 
 jobs:
   nightly_check:
@@ -22,8 +26,8 @@ jobs:
     strategy:
         fail-fast: false
         matrix:
-          os: [macOS-12, ubuntu-latest]
-          python-version: ["3.10", "3.11", "3.12"]
+          os: [macOS-13, macOS-latest, ubuntu-latest]
+          python-version: ["3.11", "3.12"]
 
     steps:
     - uses: actions/checkout@v4
@@ -42,25 +46,63 @@ jobs:
           python=${{ matrix.python-version }}
           pydantic=2
 
-    - name: Install nightly pytorch-lightning and DGL
+    - name: Install nightly pytorch
       run: |
-        python -m pip install --pre dgl -f https://data.dgl.ai/wheels-test/repo.html
-        python -m pip install --pre dglgo -f https://data.dgl.ai/wheels-test/repo.html
-        python -m pip install https://github.com/Lightning-AI/lightning/archive/refs/heads/master.zip -U  
+        python -m pip install https://github.com/Lightning-AI/lightning/archive/refs/heads/master.zip -U 
 
-    - name: Install package
+    - name: Download DGL source
       run: |
-        python -m pip install . --no-deps
+        git clone --recurse-submodules https://github.com/dmlc/dgl.git
+
+    - name: Set up DGL
+      if: matrix.os != 'ubuntu-latest'
+      run: |
+        # from https://docs.dgl.ai/en/latest/install/index.html#macos
+        cd dgl/
+        mkdir build && cd build
+        cmake -DUSE_OPENMP=off -DUSE_LIBXSMM=OFF ..
+        make -j4
+        cd ../..
+
+    - name: Set up DGL (Ubuntu)
+      if: matrix.os == 'ubuntu-latest'
+      run: |
+        cd dgl/
+        mkdir build && cd build
+        cmake -DBUILD_TYPE=dev -DUSE_CUDA=OFF -DUSE_LIBXSMM=OFF ..
+        make
+        cd ../..
+
+    - name: Install DGL Python bindings
+      run: |
+        cd dgl/python
+        python setup.py install
+        python setup.py build_ext --inplace
+        cd ../..
 
-    - name: Python information
+    - uses: ./.github/actions/include-openeye
+      with:
+        openeye-license-text: ${{ secrets.OE_LICENSE }}
+        openeye-license-file: ${{ env.OE_LICENSE }}
+
+    - name: Check toolkit installations
+      shell: bash -l -c "python -u {0}"
       run: |
-        which python
-        conda info
-        conda list
+        from openff.toolkit.utils.toolkits import OPENEYE_AVAILABLE, RDKIT_AVAILABLE
+        assert str(OPENEYE_AVAILABLE).lower() == 'true'
+        assert str(RDKIT_AVAILABLE).lower() == 'true'
 
+    - name: Check Python environment
+      run: |
+        pip list
+
     - name: Import DGL
       run: |
-        python -c "import dgl"
+        python -c "import dgl; import dgl.nn"
+
+    - name: Install package
+      run: |
+        python -m pip install . --no-deps
 
     - name: Run tests
       run: |

devtools/conda-envs/nightly.yaml

@@ -1,8 +1,8 @@
 name: openff-nagl-test
 channels:
   - openeye
-  - pytorch-nightly
   - conda-forge
+  - pytorch-nightly
   - defaults
 dependencies:
     # Base depends
@@ -18,13 +18,16 @@ dependencies:
   - pydantic <3
   - rdkit
   - scipy
+  - openeye-toolkits
+
+  # gnn
+  - pytorch
+  - torchvision
+  - torchaudio
 
   # database
   - pyarrow
 
-  # gcn
-  - pytorch
-
   # parallelism
   - dask-jobqueue
 
@@ -49,6 +52,11 @@ dependencies:
   - sqlite
   - xmltodict
 
+  # building
+  - make
+  - cmake
+  - cython
+
   # Pip-only installs
   - pip:
     - git+https://github.com/openforcefield/openff-utilities.git@main

openff/nagl/nn/_dataset.py

@@ -268,9 +268,15 @@ class _LazyDGLMoleculeDataset(Dataset):
 
     @property
     def schema(self):
+        self.get_schema()
+
+    @classmethod
+    def get_schema(cls):
         import pyarrow as pa
 
         return pa.schema([pa.field("pickled", pa.binary())])
+
+
 
     def __len__(self):
         return self.n_entries
@@ -370,15 +376,15 @@ def from_arrow_dataset(
         input_dataset = ds.dataset(path, format=format)
 
         with pa.OSFile(str(output_path), "wb") as sink:
-            with pa.ipc.new_file(sink, cls.schema) as writer:
+            with pa.ipc.new_file(sink, cls.get_schema()) as writer:
                 input_batches = input_dataset.to_batches(columns=columns)
                 for input_batch in input_batches:
                     with get_mapper_to_processes(n_processes=n_processes) as mapper:
                         pickled = list(mapper(converter, input_batch.to_pylist()))
 
                     output_batch = pa.RecordBatch.from_arrays(
                         [pa.array(pickled)],
-                        schema=cls.schema
+                        schema=cls.get_schema()
                     )
                     writer.write_batch(output_batch)