Merge branch 'main' into export-D84192560

Author: AbishekS

.github/workflows/doc-build.yaml

@@ -85,12 +85,6 @@ jobs:
           set -eux
           sudo apt-get install -y pandoc
           pip install -r docs/requirements.txt
-      - name: Start Airflow
-        run: |
-          # start airflow in background
-          airflow standalone &
-          # wait 5 seconds for airflow to start
-          sleep 5
       - name: Build
         run: |
           set -ex

docs/requirements.txt

@@ -10,4 +10,3 @@ jupytext
 ipython_genutils
 # https://github.com/jupyter/nbconvert/issues/1736
 jinja2>=3.1.4
-apache-airflow==2.10.5

docs/source/advanced.rst

@@ -153,6 +153,33 @@ resource can then be used in the following manner:
 
   test_app("gpu_x2")
 
+Alternatively, you can define custom named resources in a Python module and point
+to it using the ``TORCHX_CUSTOM_NAMED_RESOURCES`` environment variable:
+
+.. code-block:: python
+
+   # my_resources.py
+   from torchx.specs import Resource
+
+   def gpu_x8_efa() -> Resource:
+       return Resource(cpu=100, gpu=8, memMB=819200, devices={"vpc.amazonaws.com/efa": 1})
+
+   def cpu_x32() -> Resource:
+       return Resource(cpu=32, gpu=0, memMB=131072)
+
+   NAMED_RESOURCES = {
+       "gpu_x8_efa": gpu_x8_efa,
+       "cpu_x32": cpu_x32,
+   }
+
+Then set the environment variable:
+
+.. code-block:: bash
+
+   export TORCHX_CUSTOM_NAMED_RESOURCES=my_resources
+
+This allows you to use your custom resources without creating a package with entry points.
+
 
 Registering Custom Components
 -------------------------------

torchx/cli/cmd_run.py

@@ -26,7 +26,7 @@
 from torchx.runner import config, get_runner, Runner
 from torchx.runner.config import load_sections
 from torchx.schedulers import get_default_scheduler_name, get_scheduler_factories
-from torchx.specs import CfgVal
+from torchx.specs import CfgVal, Workspace
 from torchx.specs.finder import (
     _Component,
     ComponentNotFoundException,
@@ -36,7 +36,6 @@
 )
 from torchx.util.log_tee_helpers import tee_logs
 from torchx.util.types import none_throws
-from torchx.workspace import Workspace
 
 
 MISSING_COMPONENT_ERROR_MSG = (

torchx/runner/api.py

@@ -1,4 +1,3 @@
-#!/usr/bin/env python3
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -43,6 +42,7 @@
     parse_app_handle,
     runopts,
     UnknownAppException,
+    Workspace,
 )
 from torchx.specs.finder import get_component
 from torchx.tracker.api import (
@@ -54,7 +54,7 @@
 from torchx.util.session import get_session_id_or_create_new, TORCHX_INTERNAL_SESSION_ID
 
 from torchx.util.types import none_throws
-from torchx.workspace.api import Workspace, WorkspaceMixin
+from torchx.workspace import WorkspaceMixin
 
 if TYPE_CHECKING:
     from typing_extensions import Self

torchx/runner/test/config_test.py

@@ -25,9 +25,8 @@
 )
 from torchx.schedulers import get_scheduler_factories, Scheduler
 from torchx.schedulers.api import DescribeAppResponse, ListAppResponse, Stream
-from torchx.specs import AppDef, AppDryRunInfo, CfgVal, runopts
+from torchx.specs import AppDef, AppDryRunInfo, CfgVal, runopts, Workspace
 from torchx.test.fixtures import TestWithTmpDir
-from torchx.workspace import Workspace
 
 
 class TestScheduler(Scheduler):

torchx/schedulers/api.py

@@ -1,4 +1,3 @@
-#!/usr/bin/env python3
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -22,8 +21,9 @@
     Role,
     RoleStatus,
     runopts,
+    Workspace,
 )
-from torchx.workspace.api import Workspace, WorkspaceMixin
+from torchx.workspace import WorkspaceMixin
 
 
 DAYS_IN_2_WEEKS = 14

torchx/specs/__init__.py

@@ -12,6 +12,8 @@
 scheduler or pipeline adapter.
 """
 import difflib
+
+import os
 from typing import Callable, Dict, Mapping, Optional
 
 from torchx.specs.api import (
@@ -64,8 +66,10 @@
 GENERIC_NAMED_RESOURCES: Mapping[str, ResourceFactory] = import_attr(
     "torchx.specs.named_resources_generic", "NAMED_RESOURCES", default={}
 )
-FB_NAMED_RESOURCES: Mapping[str, ResourceFactory] = import_attr(
-    "torchx.specs.fb.named_resources", "NAMED_RESOURCES", default={}
+CUSTOM_NAMED_RESOURCES: Mapping[str, ResourceFactory] = import_attr(
+    os.environ.get("TORCHX_CUSTOM_NAMED_RESOURCES", "torchx.specs.fb.named_resources"),
+    "NAMED_RESOURCES",
+    default={},
 )
 
 
@@ -76,7 +80,7 @@ def _load_named_resources() -> Dict[str, Callable[[], Resource]]:
     for name, resource in {
         **GENERIC_NAMED_RESOURCES,
         **AWS_NAMED_RESOURCES,
-        **FB_NAMED_RESOURCES,
+        **CUSTOM_NAMED_RESOURCES,
         **resource_methods,
     }.items():
         materialized_resources[name] = resource

torchx/specs/named_resources_aws.py

@@ -120,6 +120,16 @@ def aws_p5_48xlarge() -> Resource:
     )
 
 
+def aws_p5e_48xlarge() -> Resource:
+    return Resource(
+        cpu=192,
+        gpu=8,
+        memMB=2048 * GiB,
+        capabilities={K8S_ITYPE: "p5e.48xlarge"},
+        devices={EFA_DEVICE: 32},
+    )
+
+
 def aws_p5en_48xlarge() -> Resource:
     return Resource(
         cpu=192,
@@ -419,6 +429,7 @@ def aws_inf2_48xlarge() -> Resource:
     "aws_p4d.24xlarge": aws_p4d_24xlarge,
     "aws_p4de.24xlarge": aws_p4de_24xlarge,
     "aws_p5.48xlarge": aws_p5_48xlarge,
+    "aws_p5e.48xlarge": aws_p5e_48xlarge,
     "aws_p5en.48xlarge": aws_p5en_48xlarge,
     "aws_g4dn.xlarge": aws_g4dn_xlarge,
     "aws_g4dn.2xlarge": aws_g4dn_2xlarge,

torchx/specs/test/named_resources_aws_test.py

@@ -44,6 +44,7 @@
     aws_p4d_24xlarge,
     aws_p4de_24xlarge,
     aws_p5_48xlarge,
+    aws_p5e_48xlarge,
     aws_p5en_48xlarge,
     aws_t3_medium,
     aws_trn1_2xlarge,
@@ -95,13 +96,19 @@ def test_aws_p4(self) -> None:
 
     def test_aws_p5(self) -> None:
         p5 = aws_p5_48xlarge()
+        p5e = aws_p5e_48xlarge()
         p5en = aws_p5en_48xlarge()
 
         self.assertEqual(192, p5.cpu)
         self.assertEqual(8, p5.gpu)
         self.assertEqual(2048 * GiB, p5.memMB)
         self.assertEqual({EFA_DEVICE: 32}, p5.devices)
 
+        self.assertEqual(192, p5e.cpu)
+        self.assertEqual(8, p5e.gpu)
+        self.assertEqual(2048 * GiB, p5e.memMB)
+        self.assertEqual({EFA_DEVICE: 32}, p5e.devices)
+
         self.assertEqual(192, p5en.cpu)
         self.assertEqual(8, p5en.gpu)
         self.assertEqual(2048 * GiB, p5en.memMB)