torchx-nightly 2023.10.21__py3-none-any.whl → 2025.12.24__py3-none-any.whl

torchx/components/dist.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-strict
+
 """
 For distributed training, TorchX relies on the scheduler's gang scheduling
 capabilities to schedule ``n`` copies of nodes. Once launched, the application
@@ -90,6 +92,7 @@ def spmd(
     h: str = "gpu.small",
     j: str = "1x1",
     env: Optional[Dict[str, str]] = None,
+    metadata: Optional[Dict[str, str]] = None,
     max_retries: int = 0,
     mounts: Optional[List[str]] = None,
     debug: bool = False,
@@ -129,10 +132,8 @@ def spmd(
         h: the type of host to run on (e.g. aws_p4d.24xlarge). Must be one of the registered named resources
         j: {nnodes}x{nproc_per_node}. For GPU hosts omitting nproc_per_node will infer it from the GPU count on the host
         env: environment variables to be passed to the run (e.g. ENV1=v1,ENV2=v2,ENV3=v3)
+        metadata: metadata to be passed to the scheduler (e.g. KEY1=v1,KEY2=v2,KEY3=v3)
         max_retries: the number of scheduler retries allowed
-        rdzv_port: the port on rank0's host to use for hosting the c10d store used for rendezvous.
-                   Only takes effect when running multi-node. When running single node, this parameter
-                   is ignored and a random free port is chosen.
         mounts: (for docker based runs only) mounts to mount into the worker environment/container
                 (ex. type=<bind/volume>,src=/host,dst=/job[,readonly]).
         debug: whether to run with preset debug flags enabled
@@ -151,6 +152,7 @@ def spmd(
         h=h,
         j=str(StructuredJArgument.parse_from(h, j)),
         env=env,
+        metadata=metadata,
         max_retries=max_retries,
         mounts=mounts,
         debug=debug,
@@ -169,11 +171,14 @@ def ddp(
     memMB: int = 1024,
     j: str = "1x2",
     env: Optional[Dict[str, str]] = None,
+    metadata: Optional[Dict[str, str]] = None,
     max_retries: int = 0,
     rdzv_port: int = 29500,
     rdzv_backend: str = "c10d",
+    rdzv_conf: Optional[str] = None,
     mounts: Optional[List[str]] = None,
     debug: bool = False,
+    tee: int = 3,
 ) -> specs.AppDef:
     """
     Distributed data parallel style application (one role, multi-replica).
@@ -185,7 +190,7 @@ def ddp(
 
     Note: (cpu, gpu, memMB) parameters are mutually exclusive with ``h`` (named resource) where
           ``h`` takes precedence if specified for setting resource requirements.
-          See `registering named resources <https://pytorch.org/torchx/latest/advanced.html#registering-named-resources>`_.
+          See `registering named resources <https://meta-pytorch.org/torchx/latest/advanced.html#registering-named-resources>`_.
 
     Args:
         script_args: arguments to the main module
@@ -200,14 +205,17 @@ def ddp(
         h: a registered named resource (if specified takes precedence over cpu, gpu, memMB)
         j: [{min_nnodes}:]{nnodes}x{nproc_per_node}, for gpu hosts, nproc_per_node must not exceed num gpus
         env: environment varibles to be passed to the run (e.g. ENV1=v1,ENV2=v2,ENV3=v3)
+        metadata: metadata to be passed to the scheduler (e.g. KEY1=v1,KEY2=v2,KEY3=v3)
         max_retries: the number of scheduler retries allowed
         rdzv_port: the port on rank0's host to use for hosting the c10d store used for rendezvous.
                    Only takes effect when running multi-node. When running single node, this parameter
                    is ignored and a random free port is chosen.
         rdzv_backend: the rendezvous backend to use. Only takes effect when running multi-node.
+        rdzv_conf: the additional rendezvous configuration to use (ex. join_timeout=600,close_timeout=600,timeout=600).
         mounts: mounts to mount into the worker environment/container (ex. type=<bind/volume>,src=/host,dst=/job[,readonly]).
                 See scheduler documentation for more info.
         debug: whether to run with preset debug flags enabled
+        tee: tees the specified std stream(s) to console + file. 0: none, 1: stdout, 2: stderr, 3: both
     """
 
     if (script is None) == (m is None):
@@ -234,8 +242,8 @@ def ddp(
         # use $$ in the prefix to escape the '$' literal (rather than a string Template substitution argument)
         rdzv_endpoint = _noquote(f"$${{{macros.rank0_env}:=localhost}}:{rdzv_port}")
 
-    if env is None:
-        env = {}
+    env = env or {}
+    metadata = metadata or {}
 
     argname = StructuredNameArgument.parse_from(
         name=name,
@@ -244,6 +252,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:
@@ -253,6 +262,7 @@ def ddp(
         "torchrun",
         "--rdzv_backend",
         rdzv_backend,
+        *(["--rdzv_conf", rdzv_conf] if rdzv_conf is not None else []),
         "--rdzv_endpoint",
         rdzv_endpoint,
         "--rdzv_id",
@@ -262,7 +272,7 @@ def ddp(
         "--nproc_per_node",
         str(nproc_per_node),
         "--tee",
-        "3",
+        str(tee),
         "--role",
         "",
     ]
@@ -294,6 +304,7 @@ def ddp(
                 mounts=specs.parse_mounts(mounts) if mounts else [],
             )
         ],
+        metadata=metadata,
     )
 
 

torchx/components/integration_tests/component_provider.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-strict
+
 import os
 import tempfile
 from abc import ABC, abstractmethod
@@ -107,7 +109,7 @@ class CopyComponentProvider(ComponentProvider):
         self._dst_path = "<None>"
 
     def setUp(self) -> None:
-        if self._scheduler in ["local_cwd", "ray"]:
+        if self._scheduler in ["local_cwd"]:
             fname = "torchx_copy_test.txt"
             self._src_path: str = os.path.join(tempfile.gettempdir(), fname)
             self._dst_path: str = os.path.join(tempfile.gettempdir(), f"{fname}.copy")
@@ -124,7 +126,7 @@ class CopyComponentProvider(ComponentProvider):
     def tearDown(self) -> None:
         if os.path.exists(self._dst_path):
             os.remove(self._dst_path)
-        if self._scheduler in ["local_cwd", "ray"] and os.path.exists(self._dst_path):
+        if self._scheduler in ["local_cwd"] and os.path.exists(self._dst_path):
             os.remove(self._dst_path)
 
     def get_app_def(self) -> AppDef:

torchx/components/integration_tests/integ_tests.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-strict
+
 import inspect
 import logging
 import sys

torchx/components/serve.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-strict
+
 """
 These components aim to make it easier to interact with inference and serving
 tools such as `torchserve <https://pytorch.org/serve/>`_.

torchx/components/structured_arg.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-strict
+
 """
 Defines methods for structured (higher order) component argument parsing.
 Use the functionalities defined in this module to author components
@@ -28,8 +30,6 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional
 
-from pyre_extensions import none_throws
-
 from torchx import specs
 
 
@@ -146,7 +146,8 @@ class StructuredNameArgument:
             if m:  # use the last module name
                 run_name = m.rpartition(".")[2]
             else:  # use script name w/ no extension
-                run_name = Path(none_throws(script)).stem
+                assert script, "`script` can't be `None` here due checks above"
+                run_name = Path(script).stem
         return StructuredNameArgument(
             experiment_name or default_experiment_name, run_name
         )
@@ -191,12 +192,12 @@ class StructuredJArgument:
 
             .. doctest::
 
-                >>> str(StructuredJArgument.parse_from(h="aws_trn1.32xl", j="2"))
+                >>> str(StructuredJArgument.parse_from(h="aws_trn1.32xlarge", j="2"))
                 Traceback (most recent call last):
                     ...
-                ValueError: nproc_per_node cannot be inferred from GPU count. `trn1.32xl` is not a GPU instance. ...
+                ValueError: nproc_per_node cannot be inferred from GPU count. `trn1.32xlarge` is not a GPU instance. ...
 
-                >>> str(StructuredJArgument.parse_from(h="aws_trn1.32xl", j="2x16"))
+                >>> str(StructuredJArgument.parse_from(h="aws_trn1.32xlarge", j="2x16"))
                 '2x16'
 
         """

torchx/components/utils.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-strict
+
 """
 This contains TorchX utility components that are `ready-to-use` out of the box. These are
 components that simply execute well known binaries (e.g. ``cp``)
@@ -81,6 +83,7 @@ def sh(
     env: Optional[Dict[str, str]] = None,
     max_retries: int = 0,
     mounts: Optional[List[str]] = None,
+    entrypoint: Optional[str] = None,
 ) -> specs.AppDef:
     """
     Runs the provided command via sh. Currently sh does not support
@@ -98,21 +101,29 @@ def sh(
         max_retries: the number of scheduler retries allowed
         mounts: mounts to mount into the worker environment/container (ex. type=<bind/volume>,src=/host,dst=/job[,readonly]).
                 See scheduler documentation for more info.
+        entrypoint: the entrypoint to use for the command (defaults to sh)
     """
 
-    escaped_args = " ".join(shlex.quote(arg) for arg in args)
+    escaped_args = [shlex.quote(arg) for arg in args]
     if env is None:
         env = {}
     env.setdefault("LOGLEVEL", os.getenv("LOGLEVEL", "WARNING"))
 
+    if entrypoint is not None:
+        resolved_entrypoint = entrypoint
+        resolved_args = escaped_args
+    else:
+        resolved_entrypoint = "sh"
+        resolved_args = ["-c", " ".join(escaped_args)]
+
     return specs.AppDef(
         name="sh",
         roles=[
             specs.Role(
                 name="sh",
                 image=image,
-                entrypoint="sh",
-                args=["-c", escaped_args],
+                entrypoint=resolved_entrypoint,
+                args=resolved_args,
                 num_replicas=num_replicas,
                 resource=specs.resource(cpu=cpu, gpu=gpu, memMB=memMB, h=h),
                 env=env,
@@ -143,7 +154,7 @@ def python(
 
     Note: (cpu, gpu, memMB) parameters are mutually exclusive with ``h`` (named resource) where
           ``h`` takes precedence if specified for setting resource requirements.
-          See `registering named resources <https://pytorch.org/torchx/latest/advanced.html#registering-named-resources>`_.
+          See `registering named resources <https://meta-pytorch.org/torchx/latest/advanced.html#registering-named-resources>`_.
 
     Args:
         args: arguments passed to the program in sys.argv[1:] (ignored with `--c`)

torchx/distributed/__init__.py

@@ -48,7 +48,7 @@ def local_rank() -> int:
                 " but the `LOCAL_RANK` environment variable is not set. Will trivially return 0 for local_rank.\n"
                 " It is recommended to use torchrun/torchx to run your script or set the `LOCAL_RANK` manually.\n"
                 " For additional details see:\n"
-                "  1) https://pytorch.org/torchx/latest/components/distributed.html\n"
+                "  1) https://meta-pytorch.org/torchx/latest/components/distributed.html\n"
                 "  2) https://pytorch.org/docs/stable/elastic/run.html\n"
                 "=============================================================================================="
             )
@@ -83,9 +83,7 @@ def local_device() -> torch.device:
     if dist.is_initialized():
         default_pg = _get_default_group()
         return (
-            local_cuda_device()
-            if default_pg.options.backend == "nccl"
-            else torch.device("cpu")
+            local_cuda_device() if default_pg.name() == "nccl" else torch.device("cpu")
         )
     else:
         return torch.device("cuda") if has_cuda_devices() else torch.device("cpu")

torchx/examples/apps/datapreproc/datapreproc.py

@@ -5,6 +5,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-strict
+
 """
 Data Preprocessing App Example
 ====================================

torchx/examples/apps/lightning/data.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-strict
+
 """
 Trainer Datasets Example
 ========================
@@ -62,17 +64,17 @@ class ImageFolderSamplesDataset(datasets.ImageFolder):
 # our trainer and other components that need to load data.
 
 
-# pyre-fixme[13]: Attribute `test_ds` is never initialized.
-# pyre-fixme[13]: Attribute `train_ds` is never initialized.
-# pyre-fixme[13]: Attribute `val_ds` is never initialized.
 class TinyImageNetDataModule(pl.LightningDataModule):
     """
     TinyImageNetDataModule is a pytorch LightningDataModule for the tiny
     imagenet dataset.
     """
 
+    # pyre-fixme[13]: Attribute `test_ds` is never initialized.
     train_ds: ImageFolderSamplesDataset
+    # pyre-fixme[13]: Attribute `train_ds` is never initialized.
     val_ds: ImageFolderSamplesDataset
+    # pyre-fixme[13]: Attribute `val_ds` is never initialized.
     test_ds: ImageFolderSamplesDataset
 
     def __init__(

torchx/examples/apps/lightning/model.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-strict
+
 """
 Tiny ImageNet Model
 ====================
@@ -21,7 +23,7 @@ import pytorch_lightning as pl
 import torch
 import torch.jit
 from torch.nn import functional as F
-from torchmetrics import Accuracy
+from torchmetrics.classification import MulticlassAccuracy
 from torchvision.models.resnet import BasicBlock, ResNet
 
 
@@ -42,13 +44,12 @@ class TinyImageNetModel(pl.LightningModule):
 
         # We use the torchvision resnet model with some small tweaks to match
         # TinyImageNet.
-        m = ResNet(BasicBlock, layer_sizes)
+        m = ResNet(BasicBlock, layer_sizes, num_classes=200)
         m.avgpool = torch.nn.AdaptiveAvgPool2d(1)
-        m.fc.out_features = 200
         self.model: ResNet = m
 
-        self.train_acc = Accuracy()
-        self.val_acc = Accuracy()
+        self.train_acc = MulticlassAccuracy(num_classes=m.fc.out_features)
+        self.val_acc = MulticlassAccuracy(num_classes=m.fc.out_features)
 
     # pyre-fixme[14]
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -69,7 +70,7 @@ class TinyImageNetModel(pl.LightningModule):
     def _step(
         self,
         step_name: str,
-        acc_metric: Accuracy,
+        acc_metric: MulticlassAccuracy,
         batch: Tuple[torch.Tensor, torch.Tensor],
         batch_idx: int,
     ) -> torch.Tensor:

torchx/examples/apps/lightning/profiler.py

@@ -5,6 +5,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-strict
+
 """
 Simple Logging Profiler
 ===========================
@@ -17,18 +19,19 @@ output is used for HPO optimization with Ax.
 import time
 from typing import Dict
 
-from pytorch_lightning.loggers.base import LightningLoggerBase
-from pytorch_lightning.profiler.base import BaseProfiler
+from pytorch_lightning.loggers.logger import Logger
+
+from pytorch_lightning.profilers.profiler import Profiler
 
 
-class SimpleLoggingProfiler(BaseProfiler):
+class SimpleLoggingProfiler(Profiler):
     """
     This profiler records the duration of actions (in seconds) and reports the
     mean duration of each action to the specified logger. Reported metrics are
     in the format `duration_<event>`.
     """
 
-    def __init__(self, logger: LightningLoggerBase) -> None:
+    def __init__(self, logger: Logger) -> None:
         super().__init__()
 
         self.current_actions: Dict[str, float] = {}

torchx/examples/apps/lightning/train.py

@@ -5,6 +5,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-strict
+
 """
 Trainer Example
 =============================================
@@ -58,6 +60,7 @@ import pytorch_lightning as pl
 import torch
 from pytorch_lightning.callbacks import ModelCheckpoint
 from pytorch_lightning.loggers import TensorBoardLogger
+from torch.distributed.elastic.multiprocessing import errors
 from torchx.examples.apps.lightning.data import (
     create_random_data,
     download_data,
@@ -83,7 +86,12 @@ def parse_args(argv: List[str]) -> argparse.Namespace:
     parser.add_argument(
         "--batch_size", type=int, default=32, help="batch size to use for training"
     )
-    parser.add_argument("--num_samples", type=int, default=10, help="num_samples")
+    parser.add_argument(
+        "--num_samples",
+        type=int,
+        default=32,
+        help="number of samples in the dataset",
+    )
     parser.add_argument(
         "--data_path",
         type=str,
@@ -124,6 +132,7 @@ def get_model_checkpoint(args: argparse.Namespace) -> Optional[ModelCheckpoint]:
     )
 
 
+@errors.record
 def main(argv: List[str]) -> None:
     with tempfile.TemporaryDirectory() as tmpdir:
         args = parse_args(argv)
@@ -136,7 +145,7 @@ def main(argv: List[str]) -> None:
         if not args.data_path:
             data_path = os.path.join(tmpdir, "data")
             os.makedirs(data_path)
-            create_random_data(data_path)
+            create_random_data(data_path, args.num_samples)
         else:
             data_path = download_data(args.data_path, tmpdir)
 

torchx/examples/torchx_out_of_sync_training.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python3
+# 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.
+
+# pyre-strict
+
+# This file is for lex team to ramp up on torchX OSS. It's for training purpose.
+# Please DO NOT modify this file unless you know what you are doing.

torchx/notebook.py

@@ -5,6 +5,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-strict
+
 """
 This contains TorchX utilities for creating and running components and apps from
 an Jupyter/IPython Notebook.

torchx/runner/__init__.py

@@ -5,4 +5,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-strict
+
 from torchx.runner.api import get_runner, Runner  # noqa: F401 F403