From 3479f021e8287a1e0d353b34c24e3af8e006fd99 Mon Sep 17 00:00:00 2001
From: Kiuk Chung <kiuk@meta.com>
Date: Thu, 27 Mar 2025 22:32:09 -0700
Subject: [PATCH] Fix broken docs-build by pinning protobuf-3.20.x (#1033)

Summary:
Pull Request resolved: https://github.com/pytorch/torchx/pull/1033

Pin protobuf-3.20.x (required by kfp-1.8.22) in dev-requirements and first install docs/requirements.txt then downgrade protobuf to 3.20.x.

What we oughta do is to upgrade to `kfp-2.x` but that's a much more intrusive change as we are bumping major versions. So pin to `apache-airflow==2.6.3` so that protobuf doesn't conflict with kfp.

Reviewed By: tonykao8080

Differential Revision: D72016216
---
 .github/workflows/doc-build.yaml             | 5 ++---
 .github/workflows/pyre.yaml                  | 2 +-
 dev-requirements.txt                         | 2 ++
 torchx/schedulers/test/ray_scheduler_test.py | 7 ++++++-
 4 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/doc-build.yaml b/.github/workflows/doc-build.yaml
index aa686ba85..699213cb2 100644
--- a/.github/workflows/doc-build.yaml
+++ b/.github/workflows/doc-build.yaml
@@ -20,16 +20,15 @@ jobs:
           architecture: x64
       - name: Checkout TorchX
         uses: actions/checkout@v2
-      - name: Install Dependencies
+      - name: Install Doc Dependencies
         run: |
           set -eux
           sudo apt-get update
           sudo apt-get install -y pandoc
-          pip install -e .[dev]
           pip install -r docs/requirements.txt
       - name: Install TorchX
         run: |
-          python setup.py develop
+          pip install -e .[dev]
       - name: Start Airflow
         run: |
           # start airflow in background
diff --git a/.github/workflows/pyre.yaml b/.github/workflows/pyre.yaml
index 823ac5936..08bfe607c 100644
--- a/.github/workflows/pyre.yaml
+++ b/.github/workflows/pyre.yaml
@@ -22,6 +22,6 @@ jobs:
           set -eux
           pip install -e .[dev]
       - name: Init Lint Runner
-          lintrunner init
+        run: lintrunner init
       - name: Run Pyre
         run: scripts/pyre.sh
diff --git a/dev-requirements.txt b/dev-requirements.txt
index 987e967f0..abebbaa65 100644
--- a/dev-requirements.txt
+++ b/dev-requirements.txt
@@ -14,6 +14,8 @@ google-cloud-runtimeconfig==0.34.0
 hydra-core
 ipython
 kfp==1.8.22
+# pin protobuf to the version that is required by kfp
+protobuf==3.20.3
 mlflow-skinny
 moto~=5.0.8
 pyre-extensions
diff --git a/torchx/schedulers/test/ray_scheduler_test.py b/torchx/schedulers/test/ray_scheduler_test.py
index 230f7a97d..8fedd4773 100644
--- a/torchx/schedulers/test/ray_scheduler_test.py
+++ b/torchx/schedulers/test/ray_scheduler_test.py
@@ -390,7 +390,12 @@ def test_nonmatching_address(self) -> None:
             ):
                 _scheduler_with_client.submit(app=app, cfg={})
 
-        def _assertDictContainsSubset(self, expected, actual, msg=None):
+        def _assertDictContainsSubset(
+            self,
+            expected: dict[str, Any],
+            actual: dict[str, Any],
+            msg: Optional[str] = None,
+        ) -> None:
             # NB: implement unittest.TestCase.assertDictContainsSubsetNew() since it was removed in python-3.11
             for key, value in expected.items():
                 self.assertIn(key, actual, msg)