import Workspace from torchx.specs not torchx.workspace (#1144)

Summary:

`Workspace` class moved from `torchx.workspace` to `torchx.specs` to accomodate the new `Role.workspace` attribute in #1139.

This diff removes the re-export of `Workspace` from `torchx.workspace` and changes import-sites to correctly import `Workspace` from `torchx.specs`

Differential Revision: D84628206

Author: kiukchung

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/workspace/__init__.py

@@ -22,4 +22,4 @@
     * ``memory://foo-bar/`` an in-memory workspace for notebook/programmatic usage
 """
 
-from torchx.workspace.api import walk_workspace, Workspace, WorkspaceMixin  # noqa: F401
+from torchx.workspace.api import walk_workspace, WorkspaceMixin  # noqa: F401