Use TorchX main from GitHub with dependency compatibility fixes

nemo_run/run/logs.py

@@ -21,12 +21,17 @@
 from typing import Optional, TextIO
 
 from torchx.cli.cmd_log import _prefix_line, find_role_replicas, validate
-from torchx.cli.colors import ENDC, GREEN
 from torchx.schedulers.api import Stream
 from torchx.specs.api import is_started, parse_app_handle
 from torchx.specs.builders import make_app_handle
 from torchx.util.types import none_throws
 
+# TorchX moved color constants from `torchx.cli.colors` to `torchx.util.colors` on main.
+try:
+    from torchx.util.colors import ENDC, GREEN
+except ImportError:  # pragma: no cover - compatibility with older TorchX releases
+    from torchx.cli.colors import ENDC, GREEN
+
 from nemo_run.core.execution.base import LogSupportedExecutor
 from nemo_run.core.frontend.console.api import CONSOLE
 from nemo_run.run.torchx_backend.runner import Runner, get_runner

nemo_run/run/torchx_backend/runner.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import inspect
 import logging
 from typing import Any, Optional
 
@@ -63,7 +64,12 @@ def dryrun(  # type: ignore
         with log_event("dryrun", scheduler):
             sched = self._scheduler(scheduler)
 
-            sched._validate(app, scheduler)
+            # TorchX `main` added `cfg` to `_validate`; older versions only accept
+            # `(app, scheduler)`. Support both signatures.
+            if "cfg" in inspect.signature(sched._validate).parameters:
+                sched._validate(app, scheduler, cfg)
+            else:
+                sched._validate(app, scheduler)
             dryrun_info = sched.submit_dryrun(app, cfg)
             dryrun_info._scheduler = scheduler
             return dryrun_info

pyproject.toml

@@ -23,16 +23,18 @@ dependencies = [
     "catalogue>=2.0.10",
     "fabric>=3.2.2",
     "fiddle>=0.3.0",
-    "torchx>=0.7.0",
+    "torchx @ git+https://github.com/meta-pytorch/torchx.git@main",
     "typer>=0.12.3",
     "rich>=13.7.1",
     "jinja2>=3.1.4",
-    "cryptography < 43.0.0",
+    "importlib-metadata",
+    "cryptography",
     "networkx >= 3.3",
     "omegaconf>=2.3.0",
     "leptonai>=0.26.6",
     "packaging",
     "toml",
+    "urllib3>=2.6.3",
 ]
 readme = "README.md"
 requires-python = ">= 3.10"