From 744571e6c0046a69cad63bc40a92fff077649c0d Mon Sep 17 00:00:00 2001
From: manav-a <33183958+manav-a@users.noreply.github.com>
Date: Fri, 5 Jul 2024 14:42:23 -0700
Subject: [PATCH] Fix up targets for devices.py and specs folder

Differential Revision: D59403775

Pull Request resolved: https://github.com/pytorch/torchx/pull/930
---
 torchx/specs/__init__.py | 55 ++++++++++++++++++++++++++++++++++------
 1 file changed, 47 insertions(+), 8 deletions(-)

diff --git a/torchx/specs/__init__.py b/torchx/specs/__init__.py
index 78aa1676e..c43cfa0c9 100644
--- a/torchx/specs/__init__.py
+++ b/torchx/specs/__init__.py
@@ -15,13 +15,7 @@
 import difflib
 from typing import Callable, Dict, Optional
 
-from torchx.specs.named_resources_aws import NAMED_RESOURCES as AWS_NAMED_RESOURCES
-from torchx.specs.named_resources_generic import (
-    NAMED_RESOURCES as GENERIC_NAMED_RESOURCES,
-)
-from torchx.util.entrypoints import load_group
-
-from .api import (  # noqa: F401 F403
+from torchx.specs.api import (
     ALL,
     AppDef,
     AppDryRunInfo,
@@ -52,7 +46,13 @@
     UnknownSchedulerException,
     VolumeMount,
 )
-from .builders import make_app_handle, materialize_appdef, parse_mounts  # noqa
+from torchx.specs.builders import make_app_handle, materialize_appdef, parse_mounts
+
+from torchx.specs.named_resources_aws import NAMED_RESOURCES as AWS_NAMED_RESOURCES
+from torchx.specs.named_resources_generic import (
+    NAMED_RESOURCES as GENERIC_NAMED_RESOURCES,
+)
+from torchx.util.entrypoints import load_group
 
 GiB: int = 1024
 
@@ -183,3 +183,42 @@ def gpu_x_1() -> Dict[str, Resource]:
 
     """
     return named_resources[res]
+
+
+__all__ = [
+    "AppDef",
+    "AppDryRunInfo",
+    "AppHandle",
+    "AppState",
+    "AppStatus",
+    "BindMount",
+    "CfgVal",
+    "DeviceMount",
+    "get_type_name",
+    "is_terminal",
+    "macros",
+    "MISSING",
+    "NONE",
+    "NULL_RESOURCE",
+    "parse_app_handle",
+    "ReplicaState",
+    "ReplicaStatus",
+    "Resource",
+    "RetryPolicy",
+    "Role",
+    "RoleStatus",
+    "runopt",
+    "runopts",
+    "UnknownAppException",
+    "UnknownSchedulerException",
+    "InvalidRunConfigException",
+    "MalformedAppHandleException",
+    "VolumeMount",
+    "resource",
+    "get_named_resources",
+    "named_resources",
+    "make_app_handle",
+    "materialize_appdef",
+    "parse_mounts",
+    "ALL",
+]