Reorganize files in repo

src/cloudai/__init__.py

@@ -12,61 +12,63 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from cloudai.installer.slurm_installer import SlurmInstaller
-from cloudai.installer.standalone_installer import StandaloneInstaller
-from cloudai.parser.system_parser.slurm_system_parser import SlurmSystemParser
-from cloudai.parser.system_parser.standalone_system_parser import StandaloneSystemParser
-from cloudai.runner.slurm.slurm_runner import SlurmRunner
-from cloudai.runner.standalone.standalone_runner import StandaloneRunner
-from cloudai.schema.core.strategy.command_gen_strategy import CommandGenStrategy
-from cloudai.schema.core.strategy.grading_strategy import GradingStrategy
-from cloudai.schema.core.strategy.install_strategy import InstallStrategy
-from cloudai.schema.core.strategy.job_id_retrieval_strategy import JobIdRetrievalStrategy
-from cloudai.schema.core.strategy.report_generation_strategy import ReportGenerationStrategy
-from cloudai.schema.system.slurm.slurm_system import SlurmSystem
-from cloudai.schema.system.standalone_system import StandaloneSystem
-from cloudai.schema.test_template.chakra_replay.grading_strategy import ChakraReplayGradingStrategy
-from cloudai.schema.test_template.chakra_replay.report_generation_strategy import ChakraReplayReportGenerationStrategy
-from cloudai.schema.test_template.chakra_replay.slurm_command_gen_strategy import ChakraReplaySlurmCommandGenStrategy
-from cloudai.schema.test_template.chakra_replay.slurm_install_strategy import ChakraReplaySlurmInstallStrategy
-from cloudai.schema.test_template.chakra_replay.template import ChakraReplay
-from cloudai.schema.test_template.common.slurm_job_id_retrieval_strategy import SlurmJobIdRetrievalStrategy
-from cloudai.schema.test_template.common.standalone_job_id_retrieval_strategy import StandaloneJobIdRetrievalStrategy
-from cloudai.schema.test_template.jax_toolbox.grading_strategy import JaxToolboxGradingStrategy
-from cloudai.schema.test_template.jax_toolbox.report_generation_strategy import JaxToolboxReportGenerationStrategy
-from cloudai.schema.test_template.jax_toolbox.slurm_command_gen_strategy import JaxToolboxSlurmCommandGenStrategy
-from cloudai.schema.test_template.jax_toolbox.slurm_install_strategy import JaxToolboxSlurmInstallStrategy
-from cloudai.schema.test_template.jax_toolbox.template import JaxToolbox
-from cloudai.schema.test_template.nccl_test.grading_strategy import NcclTestGradingStrategy
-from cloudai.schema.test_template.nccl_test.report_generation_strategy import NcclTestReportGenerationStrategy
-from cloudai.schema.test_template.nccl_test.slurm_command_gen_strategy import NcclTestSlurmCommandGenStrategy
-from cloudai.schema.test_template.nccl_test.slurm_install_strategy import NcclTestSlurmInstallStrategy
-from cloudai.schema.test_template.nccl_test.template import NcclTest
-from cloudai.schema.test_template.nemo_launcher.grading_strategy import NeMoLauncherGradingStrategy
-from cloudai.schema.test_template.nemo_launcher.report_generation_strategy import NeMoLauncherReportGenerationStrategy
-from cloudai.schema.test_template.nemo_launcher.slurm_command_gen_strategy import NeMoLauncherSlurmCommandGenStrategy
-from cloudai.schema.test_template.nemo_launcher.slurm_install_strategy import NeMoLauncherSlurmInstallStrategy
-from cloudai.schema.test_template.nemo_launcher.slurm_job_id_retrieval_strategy import (
-    NeMoLauncherSlurmJobIdRetrievalStrategy,
-)
-from cloudai.schema.test_template.nemo_launcher.template import NeMoLauncher
-from cloudai.schema.test_template.sleep.grading_strategy import SleepGradingStrategy
-from cloudai.schema.test_template.sleep.report_generation_strategy import SleepReportGenerationStrategy
-from cloudai.schema.test_template.sleep.standalone_command_gen_strategy import SleepStandaloneCommandGenStrategy
-from cloudai.schema.test_template.sleep.standalone_install_strategy import SleepStandaloneInstallStrategy
-from cloudai.schema.test_template.sleep.template import Sleep
-from cloudai.schema.test_template.ucc_test.grading_strategy import UCCTestGradingStrategy
-from cloudai.schema.test_template.ucc_test.report_generation_strategy import UCCTestReportGenerationStrategy
-from cloudai.schema.test_template.ucc_test.slurm_command_gen_strategy import UCCTestSlurmCommandGenStrategy
-from cloudai.schema.test_template.ucc_test.slurm_install_strategy import UCCTestSlurmInstallStrategy
-from cloudai.schema.test_template.ucc_test.template import UCCTest
-
+from ._core.command_gen_strategy import CommandGenStrategy
+from ._core.grading_strategy import GradingStrategy
+from ._core.install_strategy import InstallStrategy
+from ._core.job_id_retrieval_strategy import JobIdRetrievalStrategy
+from ._core.parser import Parser
 from ._core.registry import Registry
+from ._core.report_generation_strategy import ReportGenerationStrategy
+from ._core.runner import Runner
+from ._core.system import System
+from ._core.test_template import TestTemplate
+from ._core.test_template_strategy import TestTemplateStrategy
 from .grader import Grader
 from .installer.installer import Installer
-from .parser.core.parser import Parser
+from .installer.slurm_installer import SlurmInstaller
+from .installer.standalone_installer import StandaloneInstaller
+from .parser.system_parser.slurm_system_parser import SlurmSystemParser
+from .parser.system_parser.standalone_system_parser import StandaloneSystemParser
 from .report_generator import ReportGenerator
-from .runner.core.runner import Runner
+from .runner.slurm.slurm_runner import SlurmRunner
+from .runner.standalone.standalone_runner import StandaloneRunner
+from .schema.system.slurm.slurm_system import SlurmSystem
+from .schema.system.standalone_system import StandaloneSystem
+from .schema.test_template.chakra_replay.grading_strategy import ChakraReplayGradingStrategy
+from .schema.test_template.chakra_replay.report_generation_strategy import ChakraReplayReportGenerationStrategy
+from .schema.test_template.chakra_replay.slurm_command_gen_strategy import ChakraReplaySlurmCommandGenStrategy
+from .schema.test_template.chakra_replay.slurm_install_strategy import ChakraReplaySlurmInstallStrategy
+from .schema.test_template.chakra_replay.template import ChakraReplay
+from .schema.test_template.common.slurm_job_id_retrieval_strategy import SlurmJobIdRetrievalStrategy
+from .schema.test_template.common.standalone_job_id_retrieval_strategy import StandaloneJobIdRetrievalStrategy
+from .schema.test_template.jax_toolbox.grading_strategy import JaxToolboxGradingStrategy
+from .schema.test_template.jax_toolbox.report_generation_strategy import JaxToolboxReportGenerationStrategy
+from .schema.test_template.jax_toolbox.slurm_command_gen_strategy import JaxToolboxSlurmCommandGenStrategy
+from .schema.test_template.jax_toolbox.slurm_install_strategy import JaxToolboxSlurmInstallStrategy
+from .schema.test_template.jax_toolbox.template import JaxToolbox
+from .schema.test_template.nccl_test.grading_strategy import NcclTestGradingStrategy
+from .schema.test_template.nccl_test.report_generation_strategy import NcclTestReportGenerationStrategy
+from .schema.test_template.nccl_test.slurm_command_gen_strategy import NcclTestSlurmCommandGenStrategy
+from .schema.test_template.nccl_test.slurm_install_strategy import NcclTestSlurmInstallStrategy
+from .schema.test_template.nccl_test.template import NcclTest
+from .schema.test_template.nemo_launcher.grading_strategy import NeMoLauncherGradingStrategy
+from .schema.test_template.nemo_launcher.report_generation_strategy import NeMoLauncherReportGenerationStrategy
+from .schema.test_template.nemo_launcher.slurm_command_gen_strategy import NeMoLauncherSlurmCommandGenStrategy
+from .schema.test_template.nemo_launcher.slurm_install_strategy import NeMoLauncherSlurmInstallStrategy
+from .schema.test_template.nemo_launcher.slurm_job_id_retrieval_strategy import (
+    NeMoLauncherSlurmJobIdRetrievalStrategy,
+)
+from .schema.test_template.nemo_launcher.template import NeMoLauncher
+from .schema.test_template.sleep.grading_strategy import SleepGradingStrategy
+from .schema.test_template.sleep.report_generation_strategy import SleepReportGenerationStrategy
+from .schema.test_template.sleep.standalone_command_gen_strategy import SleepStandaloneCommandGenStrategy
+from .schema.test_template.sleep.standalone_install_strategy import SleepStandaloneInstallStrategy
+from .schema.test_template.sleep.template import Sleep
+from .schema.test_template.ucc_test.grading_strategy import UCCTestGradingStrategy
+from .schema.test_template.ucc_test.report_generation_strategy import UCCTestReportGenerationStrategy
+from .schema.test_template.ucc_test.slurm_command_gen_strategy import UCCTestSlurmCommandGenStrategy
+from .schema.test_template.ucc_test.slurm_install_strategy import UCCTestSlurmInstallStrategy
+from .schema.test_template.ucc_test.template import UCCTest
 
 Registry().add_system_parser("standalone", StandaloneSystemParser)
 Registry().add_system_parser("slurm", SlurmSystemParser)
@@ -122,4 +124,7 @@
     "Parser",
     "ReportGenerator",
     "Runner",
+    "System",
+    "TestTemplate",
+    "TestTemplateStrategy",
 ]

src/cloudai/runner/core/base_job.py → src/cloudai/_core/base_job.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from cloudai.schema.core.test import Test
+from .test import Test
 
 
 class BaseJob:

src/cloudai/parser/core/base_multi_file_parser.py → src/cloudai/_core/base_multi_file_parser.py

@@ -17,8 +17,9 @@
 from typing import Any, Dict, List, Union
 
 import toml
-from cloudai.schema.core.test import Test
-from cloudai.schema.core.test_template import TestTemplate
+
+from .test import Test
+from .test_template import TestTemplate
 
 
 class BaseMultiFileParser(ABC):

src/cloudai/runner/core/base_runner.py → src/cloudai/_core/base_runner.py

@@ -23,11 +23,10 @@
 from types import FrameType
 from typing import Dict, List, Optional
 
-from cloudai.schema.core.system import System
-from cloudai.schema.core.test import Test
-from cloudai.schema.core.test_scenario import TestScenario
-
 from .base_job import BaseJob
+from .system import System
+from .test import Test
+from .test_scenario import TestScenario
 
 
 class BaseRunner(ABC):

src/cloudai/parser/core/base_system_parser.py → src/cloudai/_core/base_system_parser.py

@@ -15,7 +15,7 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict
 
-from cloudai.schema.core.system import System
+from .system import System
 
 
 class BaseSystemParser(ABC):

src/cloudai/grader/grader.py → src/cloudai/_core/grader.py

@@ -17,8 +17,8 @@
 import os
 from typing import Dict, List
 
-from cloudai.schema.core.test import Test
-from cloudai.schema.core.test_scenario import TestScenario
+from .test import Test
+from .test_scenario import TestScenario
 
 
 class Grader:

src/cloudai/parser/core/parser.py → src/cloudai/_core/parser.py

@@ -15,13 +15,12 @@
 import logging
 from typing import List, Optional, Tuple
 
-from cloudai.schema.core.system import System
-from cloudai.schema.core.test_scenario import TestScenario
-from cloudai.schema.core.test_template import TestTemplate
-
+from .system import System
 from .system_parser import SystemParser
 from .test_parser import TestParser
+from .test_scenario import TestScenario
 from .test_scenario_parser import TestScenarioParser
+from .test_template import TestTemplate
 from .test_template_parser import TestTemplateParser
 
 

src/cloudai/_core/registry.py

@@ -1,13 +1,14 @@
 from typing import Dict, List, Tuple, Type, Union
 
 from cloudai.installer.base_installer import BaseInstaller
-from cloudai.parser.core.base_system_parser import BaseSystemParser
-from cloudai.runner.core.base_runner import BaseRunner
-from cloudai.schema.core.strategy.job_id_retrieval_strategy import JobIdRetrievalStrategy
-from cloudai.schema.core.strategy.report_generation_strategy import ReportGenerationStrategy
-from cloudai.schema.core.strategy.test_template_strategy import TestTemplateStrategy
-from cloudai.schema.core.system import System
-from cloudai.schema.core.test_template import TestTemplate
+
+from .base_runner import BaseRunner
+from .base_system_parser import BaseSystemParser
+from .job_id_retrieval_strategy import JobIdRetrievalStrategy
+from .report_generation_strategy import ReportGenerationStrategy
+from .system import System
+from .test_template import TestTemplate
+from .test_template_strategy import TestTemplateStrategy
 
 
 class Singleton(type):

src/cloudai/runner/core/runner.py → src/cloudai/_core/runner.py

@@ -14,11 +14,10 @@
 
 import logging
 
-from cloudai._core.registry import Registry
-from cloudai.schema.core.system import System
-from cloudai.schema.core.test_scenario import TestScenario
-
 from .base_runner import BaseRunner
+from .registry import Registry
+from .system import System
+from .test_scenario import TestScenario
 
 
 class Runner:

src/cloudai/schema/core/strategy/strategy_registry.py → src/cloudai/_core/strategy_registry.py

@@ -14,7 +14,7 @@
 
 from typing import Optional, Type
 
-from cloudai._core.registry import Registry
+from .registry import Registry
 
 
 class StrategyRegistry:

src/cloudai/parser/core/system_parser.py → src/cloudai/_core/system_parser.py

@@ -15,8 +15,9 @@
 import os
 
 import toml
-from cloudai._core.registry import Registry
-from cloudai.schema.core.system import System
+
+from .registry import Registry
+from .system import System
 
 
 class SystemParser:

src/cloudai/parser/core/test_parser.py → src/cloudai/_core/test_parser.py

@@ -14,10 +14,9 @@
 
 from typing import Any, Dict, List, Set
 
-from cloudai.schema.core.test import Test
-from cloudai.schema.core.test_template import TestTemplate
-
 from .base_multi_file_parser import BaseMultiFileParser
+from .test import Test
+from .test_template import TestTemplate
 
 
 class TestParser(BaseMultiFileParser):

src/cloudai/parser/core/test_scenario_parser.py → src/cloudai/_core/test_scenario_parser.py

@@ -17,9 +17,10 @@
 from typing import Any, Dict
 
 import toml
-from cloudai.schema.core.system import System
-from cloudai.schema.core.test import Test, TestDependency
-from cloudai.schema.core.test_scenario import TestScenario
+
+from .system import System
+from .test import Test, TestDependency
+from .test_scenario import TestScenario
 
 
 class TestScenarioParser:

src/cloudai/schema/core/test_template.py → src/cloudai/_core/test_template.py

@@ -15,12 +15,11 @@
 import logging
 from typing import Any, Dict, List, Optional
 
-from cloudai.schema.core.strategy.command_gen_strategy import CommandGenStrategy
-from cloudai.schema.core.strategy.grading_strategy import GradingStrategy
-from cloudai.schema.core.strategy.install_strategy import InstallStrategy
-from cloudai.schema.core.strategy.job_id_retrieval_strategy import JobIdRetrievalStrategy
-from cloudai.schema.core.strategy.report_generation_strategy import ReportGenerationStrategy
-
+from .command_gen_strategy import CommandGenStrategy
+from .grading_strategy import GradingStrategy
+from .install_strategy import InstallStrategy
+from .job_id_retrieval_strategy import JobIdRetrievalStrategy
+from .report_generation_strategy import ReportGenerationStrategy
 from .system import System
 
 

src/cloudai/parser/core/test_template_parser.py → src/cloudai/_core/test_template_parser.py

@@ -15,17 +15,16 @@
 import logging
 from typing import Any, Dict, Optional, Type, Union, cast
 
-from cloudai._core.registry import Registry
-from cloudai.schema.core.strategy.command_gen_strategy import CommandGenStrategy
-from cloudai.schema.core.strategy.grading_strategy import GradingStrategy
-from cloudai.schema.core.strategy.install_strategy import InstallStrategy
-from cloudai.schema.core.strategy.job_id_retrieval_strategy import JobIdRetrievalStrategy
-from cloudai.schema.core.strategy.report_generation_strategy import ReportGenerationStrategy
-from cloudai.schema.core.strategy.test_template_strategy import TestTemplateStrategy
-from cloudai.schema.core.system import System
-from cloudai.schema.core.test_template import TestTemplate
-
 from .base_multi_file_parser import BaseMultiFileParser
+from .command_gen_strategy import CommandGenStrategy
+from .grading_strategy import GradingStrategy
+from .install_strategy import InstallStrategy
+from .job_id_retrieval_strategy import JobIdRetrievalStrategy
+from .registry import Registry
+from .report_generation_strategy import ReportGenerationStrategy
+from .system import System
+from .test_template import TestTemplate
+from .test_template_strategy import TestTemplateStrategy
 
 logger = logging.getLogger(__name__)
 

src/cloudai/schema/core/strategy/test_template_strategy.py → src/cloudai/_core/test_template_strategy.py

@@ -16,7 +16,7 @@
 
 from cloudai.schema.system import SlurmSystem
 
-from ..system import System
+from .system import System
 
 
 class TestTemplateStrategy:
@@ -39,12 +39,7 @@ class TestTemplateStrategy:
 
     __test__ = False
 
-    def __init__(
-        self,
-        system: System,
-        env_vars: Dict[str, Any],
-        cmd_args: Dict[str, Any],
-    ) -> None:
+    def __init__(self, system: System, env_vars: Dict[str, Any], cmd_args: Dict[str, Any]) -> None:
         """
         Initialize a TestTemplateStrategy instance with system configuration, env variables, and command-line arguments.
 

src/cloudai/grader/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .grader import Grader
+from .._core.grader import Grader
 
 __all__ = [
     "Grader",

src/cloudai/installer/base_installer.py

@@ -17,8 +17,8 @@
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Iterable
 
-from cloudai.schema.core.system import System
-from cloudai.schema.core.test_template import TestTemplate
+from cloudai._core.system import System
+from cloudai._core.test_template import TestTemplate
 
 
 class BaseInstaller:

src/cloudai/installer/installer.py

@@ -16,8 +16,8 @@
 from typing import Iterable
 
 from cloudai._core.registry import Registry
-from cloudai.schema.core.system import System
-from cloudai.schema.core.test_template import TestTemplate
+from cloudai._core.system import System
+from cloudai._core.test_template import TestTemplate
 
 from .base_installer import BaseInstaller
 

src/cloudai/installer/slurm_installer.py

@@ -18,8 +18,8 @@
 from typing import Iterable, cast
 
 import toml
-from cloudai.schema.core.system import System
-from cloudai.schema.core.test_template import TestTemplate
+from cloudai._core.system import System
+from cloudai._core.test_template import TestTemplate
 from cloudai.schema.system import SlurmSystem
 
 from .base_installer import BaseInstaller

src/cloudai/parser/system_parser/slurm_system_parser.py

@@ -15,7 +15,7 @@
 import os
 from typing import Any, Dict, List
 
-from cloudai.parser.core.base_system_parser import BaseSystemParser
+from cloudai._core.base_system_parser import BaseSystemParser
 from cloudai.schema.system.slurm import SlurmNode, SlurmNodeState, SlurmSystem