add CI for StatefulDataLoader; fix generator pickling for windows/mac

.github/workflows/ci.yml

@@ -89,4 +89,5 @@ jobs:
           pytest --no-header -v test --ignore=test/test_period.py --ignore=test/test_text_examples.py
           --ignore=test/test_audio_examples.py --ignore=test/test_aistore.py
           --ignore=test/dataloader2/test_dataloader2.py --ignore=test/dataloader2/test_mprs.py
-          --ignore=test/test_distributed.py
+          --ignore=test/test_distributed.py --ignore=test/stateful_dataloader/test_dataloader.py
+          --ignore=test/stateful_dataloader/test_state_dict.py

.github/workflows/stateful_dataloader_ci.yml

@@ -0,0 +1,88 @@
+name: Run StatefulDataLoader Tests
+on:
+  push:
+    branches:
+      - main
+      - release/*
+    tags:
+  pull_request:
+    types: [opened, synchronize, reopened, labeled]
+    branches:
+      - main
+      # For PR created by ghstack
+      - gh/*/*/base
+      - release/*
+
+jobs:
+  test:
+    if:
+      ${{ github.repository_owner == 'pytorch' && (github.event.action != 'labeled' ||
+      startsWith(github.event.label.name, 'ciflow')) }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os:
+          - macos-latest
+          - ubuntu-latest
+          - windows-latest
+        python-version:
+          - 3.8
+          - 3.9
+          - "3.10"
+    steps:
+      - name: Get PyTorch Channel
+        shell: bash
+        run: |
+          if [[ "${{ github.base_ref }}" == release/* ]] || [[ "${{ github.ref }}" == refs/heads/release/* ]] || [[ "${{ github.ref }}" == refs/tags/v* ]]; then
+            PT_CHANNEL="https://download.pytorch.org/whl/test/cpu/torch_test.html"
+          else
+            PT_CHANNEL="https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html"
+          fi
+          echo "value=$PT_CHANNEL" >> $GITHUB_OUTPUT
+        id: pytorch_channel
+      - name: Setup additional system libraries
+        if: startsWith( matrix.os, 'ubuntu' )
+        run: |
+          sudo add-apt-repository multiverse
+          sudo apt update
+          sudo apt install rar unrar libssl-dev libcurl4-openssl-dev zlib1g-dev
+      - name: Setup Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Setup msbuild on Windows
+        if: matrix.os == 'windows-latest'
+        uses: microsoft/[email protected]
+      - name: Set up Visual Studio shell
+        if: matrix.os == 'windows-latest'
+        uses: egor-tensin/vs-shell@v2
+        with:
+          arch: x64
+      - name: Check out source repository
+        uses: actions/checkout@v3
+        with:
+          submodules: recursive
+      - name: Install dependencies
+        run: |
+          pip3 install -r requirements.txt
+          pip3 install networkx
+          pip3 install --pre torch -f "${{ steps.pytorch_channel.outputs.value }}"
+          pip3 install cmake ninja
+          echo "/home/runner/.local/bin" >> $GITHUB_PATH
+      - name: Build TorchData
+        run: |
+          pip3 install .
+        env:
+          BUILD_S3: 1
+      - name: Install test requirements
+        run: pip3 install -r test/requirements.txt
+      - name: Test documentation examples
+        run: |
+          cd ./docs
+          pip3 install -r requirements.txt
+          make doctest
+          cd ..
+      - name: Run StatefulDataLoader tests with pytest
+        if: ${{ ! contains(github.event.pull_request.labels.*.name, 'ciflow/slow') }}
+        run: pytest --no-header -v test/stateful_dataloader

test/stateful_dataloader/test_state_dict.py

@@ -65,6 +65,20 @@ def __init__(self, size, shuffle):
         self.g = torch.Generator()
         self.g.manual_seed(1)
 
+    def __getstate__(self):
+        """pickling generators fails on windows and mac, this makes sure
+        unit tests can proceed on those platforms
+        """
+        state = dict(self.__dict__)
+        del state["g"]
+        state["g_state"] = self.g.get_state()
+
+    def __setstate__(self, state):
+        g_state = state.pop("g_state")
+        self.__dict__ = state
+        self.g = torch.Generator()
+        self.g.set_state(g_state)
+
     def __len__(self):
         return self.size
 

torchdata/stateful_dataloader/stateful_dataloader.py

@@ -1369,6 +1369,3 @@ def _clean_up_worker(w):
 
     def __del__(self):
         self._shutdown_workers()
-
-
-torch.utils.data.DataLoader = StatefulDataLoader  # type: ignore[misc]