torchx-nightly 2024.1.24__py3-none-any.whl → 2024.1.25__py3-none-any.whl

torchx/schedulers/aws_batch_scheduler.py

@@ -170,7 +170,10 @@ def resource_from_resource_requirements(
 
 
 def _role_to_node_properties(
-    role: Role, start_idx: int, privileged: bool = False
+    role: Role,
+    start_idx: int,
+    privileged: bool = False,
+    job_role_arn: Optional[str] = None,
 ) -> Dict[str, object]:
     role.mounts += get_device_mounts(role.resource.devices)
 
@@ -245,6 +248,8 @@ def _role_to_node_properties(
         "mountPoints": mount_points,
         "volumes": volumes,
     }
+    if job_role_arn:
+        container["jobRoleArn"] = job_role_arn
     if role.num_replicas > 1:
         instance_type = instance_type_from_resource(role.resource)
         if instance_type is not None:
@@ -349,6 +354,7 @@ class AWSBatchOpts(TypedDict, total=False):
     privileged: bool
     share_id: Optional[str]
     priority: int
+    job_role_arn: Optional[str]
 
 
 class AWSBatchScheduler(DockerWorkspaceMixin, Scheduler[AWSBatchOpts]):
@@ -498,6 +504,7 @@ class AWSBatchScheduler(DockerWorkspaceMixin, Scheduler[AWSBatchOpts]):
                     role,
                     start_idx=node_idx,
                     privileged=cfg["privileged"],
+                    job_role_arn=cfg.get("job_role_arn"),
                 )
             )
             node_idx += role.num_replicas
@@ -573,6 +580,11 @@ class AWSBatchScheduler(DockerWorkspaceMixin, Scheduler[AWSBatchOpts]):
             "Higher number (between 0 and 9999) means higher priority. "
             "This will only take effect if the job queue has a scheduling policy.",
         )
+        opts.add(
+            "job_role_arn",
+            type_=str,
+            help="The Amazon Resource Name (ARN) of the IAM role that the container can assume for AWS permissions.",
+        )
         return opts
 
     def _get_job_id(self, app_id: str) -> Optional[str]:

torchx/schedulers/gcp_batch_scheduler.py

@@ -205,12 +205,14 @@ class GCPBatchScheduler(Scheduler[GCPBatchOpts]):
             if cpu <= 0:
                 cpu = 1
             MILLI = 1000
+            # pyre-fixme[8]: Attribute has type `Field`; used as `int`.
             res.cpu_milli = cpu * MILLI
             memMB = resource.memMB
             if memMB < 0:
                 raise ValueError(
                     f"memMB should to be set to a positive value, got {memMB}"
                 )
+            # pyre-fixme[8]: Attribute has type `Field`; used as `int`.
             res.memory_mib = memMB
 
             # TODO support named resources
@@ -358,11 +360,13 @@ class GCPBatchScheduler(Scheduler[GCPBatchOpts]):
             return None
 
         gpu = 0
+        # pyre-fixme[16]: `Field` has no attribute `instances`.
         if len(job.allocation_policy.instances) != 0:
             gpu_type = job.allocation_policy.instances[0].policy.machine_type
             gpu = GPU_TYPE_TO_COUNT[gpu_type]
 
         roles = {}
+        # pyre-fixme[16]: `RepeatedField` has no attribute `__iter__`.
         for tg in job.task_groups:
             env = tg.task_spec.environment.variables
             role = env["TORCHX_ROLE_NAME"]
@@ -386,6 +390,7 @@ class GCPBatchScheduler(Scheduler[GCPBatchOpts]):
         # TODO map role/replica status
         desc = DescribeAppResponse(
             app_id=app_id,
+            # pyre-fixme[16]: `Field` has no attribute `state`.
             state=JOB_STATE[job.status.state.name],
             roles=list(roles.values()),
         )

torchx/schedulers/kubernetes_scheduler.py

@@ -23,7 +23,7 @@ Install Volcano:
     kubectl apply -f https://raw.githubusercontent.com/volcano-sh/volcano/v1.6.0/installer/volcano-development.yaml
 
 See the
-`Volcano Quickstart <https://github.com/volcano-sh/volcano#user-content-quick-start-guide>`_
+`Volcano Quickstart <https://github.com/volcano-sh/volcano#quick-start-guide>`_
 for more information.
 """
 

torchx/tracker/__init__.py

@@ -37,7 +37,7 @@ Tracker Setup
 -------------
 To enable tracking it requires:
 
-1. Defining tracker backends (entrypoints and configuration) on launcher side using :doc:`runner.config`
+1. Defining tracker backends (entrypoints/modules and configuration) on launcher side using :doc:`runner.config`
 2. Adding entrypoints within a user job using entry_points (`specification`_)
 
 .. _specification: https://packaging.python.org/en/latest/specifications/entry-points/
@@ -49,13 +49,13 @@ To enable tracking it requires:
 User can define any number of tracker backends under **torchx:tracker** section in :doc:`runner.config`, where:
    * Key: is an arbitrary name for the tracker, where the name will be used to configure its properties
         under [tracker:<TRACKER_NAME>]
-   * Value: is *entrypoint/factory method* that must be available within user job. The value will be injected into a
+   * Value: is *entrypoint* or *module* factory method that must be available within user job. The value will be injected into a
         user job and used to construct tracker implementation.
 
 .. code-block:: ini
 
     [torchx:tracker]
-    tracker_name=<entry_point>
+    tracker_name=<entry_point_or_module_factory_method>
 
 
 Each tracker can be additionally configured (currently limited to `config` parameter) under `[tracker:<TRACKER NAME>]` section:
@@ -71,11 +71,15 @@ For example, ~/.torchxconfig may be setup as:
 
     [torchx:tracker]
     tracker1=tracker1
-    tracker12=backend_2_entry_point
+    tracker2=backend_2_entry_point
+    tracker3=torchx.tracker.mlflow:create_tracker
 
     [tracker:tracker1]
     config=s3://my_bucket/config.json
 
+    [tracker:tracker3]
+    config=my_config.json
+
 
 2. User job configuration (Advanced)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

torchx/tracker/api.py

@@ -14,6 +14,7 @@ from functools import lru_cache
 from typing import Iterable, Mapping, Optional
 
 from torchx.util.entrypoints import load_group
+from torchx.util.modules import load_module
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -177,30 +178,26 @@ def _extract_tracker_name_and_config_from_environ() -> Mapping[str, Optional[str
 
 
 def build_trackers(
-    entrypoint_and_config: Mapping[str, Optional[str]]
+    factory_and_config: Mapping[str, Optional[str]]
 ) -> Iterable[TrackerBase]:
     trackers = []
 
-    entrypoint_factories = load_group("torchx.tracker")
+    entrypoint_factories = load_group("torchx.tracker") or {}
     if not entrypoint_factories:
-        logger.warning(
-            "No 'torchx.tracker' entry_points are defined. Tracking will not capture any data."
-        )
-        return trackers
+        logger.warning("No 'torchx.tracker' entry_points are defined.")
 
-    for entrypoint_key, config in entrypoint_and_config.items():
-        if entrypoint_key not in entrypoint_factories:
+    for factory_name, config in factory_and_config.items():
+        factory = entrypoint_factories.get(factory_name) or load_module(factory_name)
+        if not factory:
             logger.warning(
-                f"Could not find `{entrypoint_key}` tracker entrypoint. Skipping..."
+                f"No tracker factory `{factory_name}` found in entry_points or modules. See https://pytorch.org/torchx/main/tracker.html#module-torchx.tracker"
             )
             continue
-        factory = entrypoint_factories[entrypoint_key]
         if config:
-            logger.info(f"Tracker config found for `{entrypoint_key}` as `{config}`")
-            tracker = factory(config)
+            logger.info(f"Tracker config found for `{factory_name}` as `{config}`")
         else:
-            logger.info(f"No tracker config specified for `{entrypoint_key}`")
-            tracker = factory(None)
+            logger.info(f"No tracker config specified for `{factory_name}`")
+        tracker = factory(config)
         trackers.append(tracker)
     return trackers
 

torchx/util/modules.py

@@ -0,0 +1,33 @@
+# 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.
+
+import importlib
+from types import ModuleType
+from typing import Callable, Optional, Union
+
+
+def load_module(path: str) -> Union[ModuleType, Optional[Callable[..., object]]]:
+    """
+    Loads and returns the module/module attr represented by the ``path``: ``full.module.path:optional_attr``
+
+    ::
+
+
+    1. ``load_module("this.is.a_module:fn")`` -> equivalent to ``this.is.a_module.fn``
+    1. ``load_module("this.is.a_module")`` -> equivalent to ``this.is.a_module``
+    """
+    parts = path.split(":", 2)
+    module_path, method = parts[0], parts[1] if len(parts) > 1 else None
+    module = None
+    i, n = -1, len(module_path)
+    try:
+        while i < n:
+            i = module_path.find(".", i + 1)
+            i = i if i >= 0 else n
+            module = importlib.import_module(module_path[:i])
+        return getattr(module, method) if method else module
+    except Exception:
+        return None

{torchx_nightly-2024.1.24.dist-info → torchx_nightly-2024.1.25.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: torchx-nightly
-Version: 2024.1.24
+Version: 2024.1.25
 Summary: TorchX SDK and Components
 Home-page: https://github.com/pytorch/torchx
 Author: TorchX Devs

{torchx_nightly-2024.1.24.dist-info → torchx_nightly-2024.1.25.dist-info}/RECORD

@@ -65,13 +65,13 @@ torchx/runtime/tracking/__init__.py,sha256=uHbJ1NqsxFWGYz2aV0_p4OCMhW467zDJu_86B
 torchx/runtime/tracking/api.py,sha256=9mlsCnnKP8hfvypNcEX2_57OYMW4AuTMM-nvsIgOzK4,5457
 torchx/schedulers/__init__.py,sha256=cCansxGU45SV_lxhgzyw2on7AJyIvhprAFo6Di1x9xQ,2157
 torchx/schedulers/api.py,sha256=XlYrD6ZjV71HotJxdVZxA_Zc8DuxhM4KKCnkibqZflU,14140
-torchx/schedulers/aws_batch_scheduler.py,sha256=H93vEYQhYMdODJb0puzHNZSCT1ZMFqlDCHdYKznDJBg,27113
+torchx/schedulers/aws_batch_scheduler.py,sha256=t6wGK1NUjhTAoxH3ie7lIMTl2cxwdixkcT3HxKsggDk,27517
 torchx/schedulers/devices.py,sha256=PNbcpf8fEM18Ag1RgK9Q30zPBalEcPdsFWctdbLxuv8,1352
 torchx/schedulers/docker_scheduler.py,sha256=I-kZN-dXoQyokLPe9ZKjfhkVX5lHx_C5jvLLc2JmXQQ,15456
-torchx/schedulers/gcp_batch_scheduler.py,sha256=pgPnlx5tzogPw6FuOkoaxBY-rVjEh3mWFP-yX5izGZ8,16206
+torchx/schedulers/gcp_batch_scheduler.py,sha256=mBxJbrNTUbIuYmudzyhOOcf8KAuUpxhiQMTDgJPtL8M,16549
 torchx/schedulers/ids.py,sha256=IGsJEbCYTdfKdU3MhKLQU6b7sWCJy5dlRV6JIL_9BlE,1783
 torchx/schedulers/kubernetes_mcad_scheduler.py,sha256=xAt-on3K8HwS2kzWasn0zXd2q4IDQzo2N5A5Ehh9NII,42885
-torchx/schedulers/kubernetes_scheduler.py,sha256=loKbNhDkR6MT_oAPoDAAIzULBTTbA8op9exABwBFFXw,27081
+torchx/schedulers/kubernetes_scheduler.py,sha256=6NXYJwiYCXNeB3ubr8t4q_SuAa-vlYdiCAPXTB3f-zg,27068
 torchx/schedulers/local_scheduler.py,sha256=QUxOLRN8CSmlEUlF_cCKveZeWcieXor2rAuAGtHNzfE,40133
 torchx/schedulers/lsf_scheduler.py,sha256=KM4-LSBiTYdtP1Js8F9dSjAdNwimaTaLraZmgnZiRuI,17638
 torchx/schedulers/ray_scheduler.py,sha256=unnDtDu1rPpCLJLDcm4NYRo9ZCCtQgG5BtlHwVfly-U,17448
@@ -93,8 +93,8 @@ torchx/specs/test/components/a/b/__init__.py,sha256=Md3cCHD7Ano9kV15PqGbicgUO-RM
 torchx/specs/test/components/a/b/c.py,sha256=QyTZfsCaSZscmk3DeNOkAyMoz6GCcayrWtOKbNFIZ1M,539
 torchx/specs/test/components/c/__init__.py,sha256=fE0IHi1JJpxsNVBNzWNee2thrNXFFRhY94c80RxNSIE,231
 torchx/specs/test/components/c/d.py,sha256=RH07jjo6uvFbzIaNFnAwmD_h24cEsT8kyZDTN-ezFio,531
-torchx/tracker/__init__.py,sha256=bnqiFZmrQoioeDlUkrzxFYzfbuf6OQvrlZmhtNekw8U,4206
-torchx/tracker/api.py,sha256=WRS3As_8YkxaL_enMz60wKba9cC_ynGIp8IQwNBbTTM,11263
+torchx/tracker/__init__.py,sha256=3uUMRTQbQ-EcJrAofn3YiYJ00QwzbK4j6bJdgX8_D-k,4350
+torchx/tracker/api.py,sha256=1K7X5mtkoEonICOEXCpPHyefcnILI_VZhqDpPb-yOmc,11254
 torchx/tracker/mlflow.py,sha256=poeoIXVPzr2sxgi515fMGRH83KAFNL6XFILMh0EQ2Dw,14487
 torchx/tracker/backend/__init__.py,sha256=fE0IHi1JJpxsNVBNzWNee2thrNXFFRhY94c80RxNSIE,231
 torchx/tracker/backend/fsspec.py,sha256=JpSioMgn54mrxqqpY0kw5Gudqx9hhxkgDLaOFSEP2Ko,10425
@@ -103,6 +103,7 @@ torchx/util/cuda.py,sha256=GiAtP9-T4zJxwHl7r5qGQRuINb59Mzwz8tnzB1MVY74,1084
 torchx/util/datetime.py,sha256=e-sO5Wjx1Utpln14C3qfJHl4v3KM-SMnn11hSyvkqFY,390
 torchx/util/entrypoints.py,sha256=C4A7cF1tPLlfyYWyZ7uZEtsKeuoOoLbMv0sOSxLhXs4,2710
 torchx/util/io.py,sha256=sxb6KI42Lq6n5z6_-YKW_mAhgPdC6CxzexlMyGheWSc,1792
+torchx/util/modules.py,sha256=PjAvkC199EYEQCRIM-Fmrb1DRnySNcnV3xWrfKtqaqQ,1116
 torchx/util/shlex.py,sha256=KzyWektMeU3oXS3Z5mFkNSPLItBTszVcvQ3EYfOMUYA,448
 torchx/util/strings.py,sha256=7CZe5WKHa7IQ6DuJCYeJ5FapUC4Fd1OGeq1yZAmjluw,663
 torchx/util/types.py,sha256=6ASuDKGO91UU3DCSuWhPX_C03341tApLCQEByUz8xpY,7016
@@ -110,9 +111,9 @@ torchx/workspace/__init__.py,sha256=KbGEzJqqXaIxALm_EQO64aw-fE7MeDMFXcpU1mY650I,
 torchx/workspace/api.py,sha256=Ej6DR__mNWaVyZgoVNAAOloDy1kTD5X1jz7pRtoVf80,5464
 torchx/workspace/dir_workspace.py,sha256=Fz-hKIx0KN8iJf2BsthNj0NvTkWlxP6WFsElPs_BaT0,2253
 torchx/workspace/docker_workspace.py,sha256=Yd8ut26bNfjyJQnmH8ANOrflfr-4VKcnOrIjbi_XIUY,9208
-torchx_nightly-2024.1.24.dist-info/LICENSE,sha256=WVHfXhFC0Ia8LTKt_nJVYobdqTJVg_4J3Crrfm2A8KQ,1721
-torchx_nightly-2024.1.24.dist-info/METADATA,sha256=lFTZxMnblxf-N0f0EU4XQS_GZBYZaXhydywm4lY5hzw,5611
-torchx_nightly-2024.1.24.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-torchx_nightly-2024.1.24.dist-info/entry_points.txt,sha256=3JYZFlX9aWzR-Gs_qsx1zq7mlqbFz6Mi9rQUULW8caI,170
-torchx_nightly-2024.1.24.dist-info/top_level.txt,sha256=pxew3bc2gsiViS0zADs0jb6kC5v8o_Yy_85fhHj_J1A,7
-torchx_nightly-2024.1.24.dist-info/RECORD,,
+torchx_nightly-2024.1.25.dist-info/LICENSE,sha256=WVHfXhFC0Ia8LTKt_nJVYobdqTJVg_4J3Crrfm2A8KQ,1721
+torchx_nightly-2024.1.25.dist-info/METADATA,sha256=ZXhFYcZqrWvvAe-gwbPkVhYDMe8uJtpqbNP29mTRDe4,5611
+torchx_nightly-2024.1.25.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+torchx_nightly-2024.1.25.dist-info/entry_points.txt,sha256=3JYZFlX9aWzR-Gs_qsx1zq7mlqbFz6Mi9rQUULW8caI,170
+torchx_nightly-2024.1.25.dist-info/top_level.txt,sha256=pxew3bc2gsiViS0zADs0jb6kC5v8o_Yy_85fhHj_J1A,7
+torchx_nightly-2024.1.25.dist-info/RECORD,,