Merge branch 'main' into patch-1

omrishiv · web-flow · commit 89f7ddaa523c · 2025-04-01T21:58:38.000+01:00
diff --git a/.github/workflows/aws-batch-integration-tests.yaml b/.github/workflows/aws-batch-integration-tests.yaml
@@ -8,7 +8,7 @@ on:
 
 jobs:
   awsbatch:
-    runs-on: "linux.20_04.4x"
+    runs-on: linux.24_04.4x
     permissions:
       id-token: write
       contents: read
diff --git a/.github/workflows/components-integration-tests.yaml b/.github/workflows/components-integration-tests.yaml
@@ -12,20 +12,20 @@ jobs:
       matrix:
         include:
           - scheduler: "aws_batch"
-            platform: "linux.20_04.4x"
+            platform: linux.24_04.4x
           - scheduler: "aws_batch"
             container_repo: localhost
             extra_args: "--mock"
-            platform: "linux.20_04.4x"
+            platform: linux.24_04.4x
           - scheduler: "kubernetes"
             container_repo: localhost:5000/torchx
-            platform: "linux.20_04.16x"
+            platform: linux.24_04.4x
           - scheduler: "local_cwd"
-            platform: ubuntu-20.04
+            platform: ubuntu-latest
           - scheduler: "local_docker"
-            platform: "linux.20_04.4x"
+            platform: linux.24_04.4x
           - scheduler: "ray"
-            platform: ubuntu-20.04
+            platform: ubuntu-latest
       fail-fast: false
     runs-on: ${{ matrix.platform }}
     permissions:
diff --git a/.github/workflows/container.yaml b/.github/workflows/container.yaml
@@ -7,7 +7,7 @@ on:
 
 jobs:
   build:
-    runs-on: ubuntu-20.04
+    runs-on: linux.24_04.4x
     permissions:
       contents: read
       packages: write
diff --git a/.github/workflows/doc-build.yaml b/.github/workflows/doc-build.yaml
@@ -11,7 +11,7 @@ permissions:
 
 jobs:
   docbuild:
-    runs-on: linux.20_04.4x
+    runs-on: linux.24_04.4x
     steps:
       - name: Setup Python
         uses: actions/setup-python@v2
@@ -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
@@ -65,7 +64,7 @@ jobs:
           fi
 
   docpush:
-    runs-on: linux.20_04.4x
+    runs-on: linux.24_04.4x
     needs: docbuild
     if: ${{ github.ref == 'refs/heads/main' }}
     steps:
diff --git a/.github/workflows/gcp-batch-integration-tests.yaml b/.github/workflows/gcp-batch-integration-tests.yaml
@@ -9,7 +9,7 @@ on:
 jobs:
   gcpbatch:
     if: github.event.pull_request.head.repo.fork == false
-    runs-on: ubuntu-20.04
+    runs-on: linux.24_04.4x
     permissions:
       id-token: write
       contents: read
diff --git a/.github/workflows/kfp-integration-tests.yaml b/.github/workflows/kfp-integration-tests.yaml
@@ -8,7 +8,7 @@ on:
 
 jobs:
   kfp-launch:
-    runs-on: linux.20_04.16x
+    runs-on: inux.24_04.16x
     steps:
       - name: Setup Python
         uses: actions/setup-python@v2
diff --git a/.github/workflows/kubernetes-minikube-integration-test.yaml b/.github/workflows/kubernetes-minikube-integration-test.yaml
@@ -8,7 +8,7 @@ on:
 
 jobs:
   kubernetes-launch:
-    runs-on: "linux.20_04.16x"
+    runs-on: linux.24_04.4x
     permissions:
       id-token: write
       contents: read
diff --git a/.github/workflows/lint.yaml b/.github/workflows/lint.yaml
@@ -8,7 +8,7 @@ on:
 
 jobs:
   lint:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-latest
     steps:
       - name: Setup Python
         uses: actions/setup-python@v2
diff --git a/.github/workflows/nightly.yaml b/.github/workflows/nightly.yaml
@@ -7,7 +7,7 @@ on:
 
 jobs:
   nightly:
-    runs-on: linux.20_04.4x
+    runs-on: linux.24_04.4x
     steps:
       - name: Setup Python 3.10
         uses: actions/setup-python@v2
diff --git a/.github/workflows/pyre.yaml b/.github/workflows/pyre.yaml
@@ -8,7 +8,7 @@ on:
 
 jobs:
   pyre:
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-latest
     steps:
       - name: Setup Python
         uses: actions/setup-python@v2
@@ -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/.github/workflows/python-unittests.yaml b/.github/workflows/python-unittests.yaml
@@ -11,10 +11,10 @@ jobs:
     strategy:
       matrix:
         python-version: [3.9, "3.10", 3.11, 3.12]
-        platform: ["linux.20_04.4x"]
+        platform: [linux.24_04.4x]
         include:
-          - python-version: 3.9
-            platform: macos-12-xl
+          - python-version: 3.12
+            platform: macos-13-xlarge
       fail-fast: false
     env:
       OS: ${{ matrix.platform }}
diff --git a/dev-requirements.txt b/dev-requirements.txt
@@ -1,6 +1,6 @@
-aiobotocore==2.12.1
+aiobotocore==2.20.0
 ax-platform[mysql]==0.2.3
-boto3==1.34.51
+boto3==1.36.0
 captum>=0.4.0
 docker
 kubernetes==25.3.0
@@ -14,15 +14,17 @@ 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
 pyre-check
 pytest
 pytest-cov
-pytorch-lightning==2.3.1
+pytorch-lightning==2.5.0
 tensorboard==2.14.0
-sagemaker==2.224.4
+sagemaker==2.230.0
 torch-model-archiver>=0.4.2
 torch>=2.6.0
 torchmetrics==1.6.3
diff --git a/docs/doc_push.sh b/docs/doc_push.sh
@@ -54,8 +54,12 @@ else
     redirects=(latest main)
 fi
 
-echo "Installing torchx from $repo_root..."
+echo "Installing TorchX and Doc dependencies from $repo_root..."
 cd "$repo_root" || exit
+
+# First install doc requirements, then install torchx
+# so that torchx's pinned requirements are honored
+pip install -r docs/requirements.txt
 pip uninstall -y torchx
 pip install -r dev-requirements.txt
 python setup.py install
@@ -66,7 +70,7 @@ echo "Building torchx-$torchx_ver docs..."
 docs_dir=$repo_root/docs
 build_dir=$docs_dir/build
 cd "$docs_dir" || exit
-pip install -r requirements.txt
+
 make clean html
 echo "Doc build complete"
 
diff --git a/fmt-requirements.txt b/fmt-requirements.txt
@@ -1,5 +1,5 @@
 # generated by `pyfmt --requirements`
-black==24.4.2
+black==24.10.0
 ruff-api==0.1.0
 stdlibs==2024.1.28
 ufmt==2.8.0
diff --git a/torchx/components/dist.py b/torchx/components/dist.py
@@ -248,6 +248,7 @@ def ddp(
     )
 
     env["TORCHX_TRACKING_EXPERIMENT_NAME"] = argname.experiment_name
+    env["TORCHX_TRACKING_RUN_NAME"] = argname.run_name
 
     env.setdefault("LOGLEVEL", os.getenv("LOGLEVEL", "WARNING"))
     if debug:
diff --git a/torchx/components/test/dist_test.py b/torchx/components/test/dist_test.py
@@ -56,27 +56,31 @@ def test_validate_spmd(self) -> None:
     def test_spmd_call_by_module_or_script_no_name(self) -> None:
         appdef = spmd(script="foo/bar.py")
         self.assertEqual("bar", appdef.name)
+        self.assertEqual("bar", appdef.roles[0].env["TORCHX_TRACKING_RUN_NAME"])
         self.assertEqual(
             "default-experiment",
             appdef.roles[0].env["TORCHX_TRACKING_EXPERIMENT_NAME"],
         )
 
         appdef = spmd("-a", "b", script="foo/bar.py")
         self.assertEqual("bar", appdef.name)
+        self.assertEqual("bar", appdef.roles[0].env["TORCHX_TRACKING_RUN_NAME"])
         self.assertEqual(
             "default-experiment",
             appdef.roles[0].env["TORCHX_TRACKING_EXPERIMENT_NAME"],
         )
 
         appdef = spmd(m="foo.bar")
         self.assertEqual("bar", appdef.name)
+        self.assertEqual("bar", appdef.roles[0].env["TORCHX_TRACKING_RUN_NAME"])
         self.assertEqual(
             "default-experiment",
             appdef.roles[0].env["TORCHX_TRACKING_EXPERIMENT_NAME"],
         )
 
         appdef = spmd("-a", "b", m="foo.bar")
         self.assertEqual("bar", appdef.name)
+        self.assertEqual("bar", appdef.roles[0].env["TORCHX_TRACKING_RUN_NAME"])
         self.assertEqual(
             "default-experiment",
             appdef.roles[0].env["TORCHX_TRACKING_EXPERIMENT_NAME"],
@@ -91,35 +95,43 @@ def test_spmd_call_by_module_or_script_no_name(self) -> None:
     def test_spmd_call_by_module_or_script_with_name(self) -> None:
         appdef = spmd(script="foo/bar.py", name="baz/trial_1")
         self.assertEqual("trial_1", appdef.name)
+        self.assertEqual("trial_1", appdef.roles[0].env["TORCHX_TRACKING_RUN_NAME"])
         self.assertEqual("baz", appdef.roles[0].env["TORCHX_TRACKING_EXPERIMENT_NAME"])
 
         appdef = spmd("-a", "b", script="foo/bar.py", name="baz/trial_1")
         self.assertEqual("trial_1", appdef.name)
         self.assertEqual("baz", appdef.roles[0].env["TORCHX_TRACKING_EXPERIMENT_NAME"])
+        self.assertEqual("trial_1", appdef.roles[0].env["TORCHX_TRACKING_RUN_NAME"])
 
         appdef = spmd(m="foo.bar", name="baz/trial_1")
         self.assertEqual("trial_1", appdef.name)
+        self.assertEqual("trial_1", appdef.roles[0].env["TORCHX_TRACKING_RUN_NAME"])
         self.assertEqual("baz", appdef.roles[0].env["TORCHX_TRACKING_EXPERIMENT_NAME"])
 
         appdef = spmd("-a", "b", m="foo.bar", name="baz/trial_1")
         self.assertEqual("trial_1", appdef.name)
+        self.assertEqual("trial_1", appdef.roles[0].env["TORCHX_TRACKING_RUN_NAME"])
         self.assertEqual("baz", appdef.roles[0].env["TORCHX_TRACKING_EXPERIMENT_NAME"])
 
     def test_spmd_call_by_module_or_script_with_experiment_name(self) -> None:
         appdef = spmd(script="foo/bar.py", name="baz/")
         self.assertEqual("bar", appdef.name)
+        self.assertEqual("bar", appdef.roles[0].env["TORCHX_TRACKING_RUN_NAME"])
         self.assertEqual("baz", appdef.roles[0].env["TORCHX_TRACKING_EXPERIMENT_NAME"])
 
         appdef = spmd("-a", "b", script="foo/bar.py", name="baz/")
         self.assertEqual("bar", appdef.name)
+        self.assertEqual("bar", appdef.roles[0].env["TORCHX_TRACKING_RUN_NAME"])
         self.assertEqual("baz", appdef.roles[0].env["TORCHX_TRACKING_EXPERIMENT_NAME"])
 
         appdef = spmd(m="foo.bar", name="baz/")
         self.assertEqual("bar", appdef.name)
+        self.assertEqual("bar", appdef.roles[0].env["TORCHX_TRACKING_RUN_NAME"])
         self.assertEqual("baz", appdef.roles[0].env["TORCHX_TRACKING_EXPERIMENT_NAME"])
 
         appdef = spmd("-a", "b", m="foo.bar", name="baz/")
         self.assertEqual("bar", appdef.name)
+        self.assertEqual("bar", appdef.roles[0].env["TORCHX_TRACKING_RUN_NAME"])
         self.assertEqual("baz", appdef.roles[0].env["TORCHX_TRACKING_EXPERIMENT_NAME"])
 
     def test_spmd_call_by_module_or_script_with_run_name(self) -> None:
@@ -129,24 +141,40 @@ def test_spmd_call_by_module_or_script_with_run_name(self) -> None:
             "default-experiment",
             appdef.roles[0].env["TORCHX_TRACKING_EXPERIMENT_NAME"],
         )
+        self.assertEqual(
+            "trial_1",
+            appdef.roles[0].env["TORCHX_TRACKING_RUN_NAME"],
+        )
 
         appdef = spmd("-a", "b", script="foo/bar.py", name="/trial_1")
         self.assertEqual("trial_1", appdef.name)
         self.assertEqual(
             "default-experiment",
             appdef.roles[0].env["TORCHX_TRACKING_EXPERIMENT_NAME"],
         )
+        self.assertEqual(
+            "trial_1",
+            appdef.roles[0].env["TORCHX_TRACKING_RUN_NAME"],
+        )
 
         appdef = spmd(m="foo.bar", name="/trial_1")
         self.assertEqual("trial_1", appdef.name)
         self.assertEqual(
             "default-experiment",
             appdef.roles[0].env["TORCHX_TRACKING_EXPERIMENT_NAME"],
         )
+        self.assertEqual(
+            "trial_1",
+            appdef.roles[0].env["TORCHX_TRACKING_RUN_NAME"],
+        )
 
         appdef = spmd("-a", "b", m="foo.bar", name="/trial_1")
         self.assertEqual("trial_1", appdef.name)
         self.assertEqual(
             "default-experiment",
             appdef.roles[0].env["TORCHX_TRACKING_EXPERIMENT_NAME"],
         )
+        self.assertEqual(
+            "trial_1",
+            appdef.roles[0].env["TORCHX_TRACKING_RUN_NAME"],
+        )
diff --git a/torchx/examples/apps/compute_world_size/config/__init__.py b/torchx/examples/apps/compute_world_size/config/__init__.py
@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
diff --git a/torchx/examples/apps/compute_world_size/module/test/util_test.py b/torchx/examples/apps/compute_world_size/module/test/util_test.py
@@ -26,8 +26,8 @@ def test_compute_world_size(self) -> None:
                     "rank": 0,
                     "world_size": 1,
                     "master_addr": "localhost",
-                    # ephemeral port range in linux
-                    "master_port": random.randint(32768, 60999),
+                    # specifying 0 as master_port makes TCPStore chose a free random port
+                    "master_port": 0,
                     "backend": "gloo",
                 }
             }
diff --git a/torchx/examples/apps/compute_world_size/module/util.py b/torchx/examples/apps/compute_world_size/module/util.py
@@ -17,10 +17,12 @@
 
 def compute_world_size(cfg: DictConfig) -> int:
     # required env vars for initializing pg with the default init_method (env://)
-    os.environ["RANK"] = str(cfg.main.rank)
-    os.environ["WORLD_SIZE"] = str(cfg.main.world_size)
-    os.environ["MASTER_ADDR"] = cfg.main.master_addr
-    os.environ["MASTER_PORT"] = str(cfg.main.master_port)
+    # read from hydra config in config/defaults.yaml if not set already
+    # this can happen is compute_world_size is run directly (not with torchrun)
+    os.environ.setdefault("RANK", str(cfg.main.rank))
+    os.environ.setdefault("WORLD_SIZE", str(cfg.main.world_size))
+    os.environ.setdefault("MASTER_ADDR", cfg.main.master_addr)
+    os.environ.setdefault("MASTER_PORT", str(cfg.main.master_port))
 
     backend = cfg.main.backend
 
diff --git a/torchx/schedulers/api.py b/torchx/schedulers/api.py
@@ -185,8 +185,8 @@ def submit_dryrun(self, app: A, cfg: T) -> D:
         if isinstance(app, AppDef):
             for role in app.roles:
                 dryrun_info = role.pre_proc(self.backend, dryrun_info)
-            dryrun_info._app = app
-            dryrun_info._cfg = resolved_cfg
+        dryrun_info._app = app
+        dryrun_info._cfg = resolved_cfg
         return dryrun_info
 
     @abc.abstractmethod
diff --git a/torchx/schedulers/test/ray_scheduler_test.py b/torchx/schedulers/test/ray_scheduler_test.py
diff --git a/torchx/specs/file_linter.py b/torchx/specs/file_linter.py
