Merge pull request #89 from TaekyungHeo/docstring-jun

TaekyungHeo · web-flow · commit 2a5598a811de · 2024-06-07T12:07:38.000-04:00
Update docstring to comply with coding style and max column length
diff --git a/src/cloudai/_core/base_system_parser.py b/src/cloudai/_core/base_system_parser.py
@@ -25,8 +25,7 @@ class BaseSystemParser(ABC):
     Parses system configuration data and creates system objects.
 
     Methods
-        parse: Abstract method to parse configuration data and return a
-               System object.
+        parse: Abstract method to parse configuration data and return a System object.
     """
 
     @abstractmethod
diff --git a/src/cloudai/_core/grader.py b/src/cloudai/_core/grader.py
@@ -27,8 +27,7 @@ class Grader:
 
     Attributes
         output_path (str): The path where the performance results are stored.
-        logger (logging.Logger): Logger for the class, used to log messages
-                                 related to the grading process.
+        logger (logging.Logger): Logger for the class, used to log messages related to the grading process.
     """
 
     def __init__(self, output_path: str) -> None:
@@ -42,8 +41,7 @@ def grade(self, test_scenario: TestScenario) -> str:
         generates a comprehensive weighted report.
 
         Args:
-            test_scenario (TestScenario): The test scenario containing multiple
-                                          tests to grade.
+            test_scenario (TestScenario): The test scenario containing multiple tests to grade.
 
         Returns:
             str: A report summarizing the weighted performance grades.
@@ -94,8 +92,7 @@ def _generate_report(self, test_perfs: Dict[str, List[float]], overall_avg: floa
         Generate a human-readable report from test performance metrics.
 
         Args:
-            test_perfs (Dict[str, List[float]]): The performance metrics for
-                                                 each test.
+            test_perfs (Dict[str, List[float]]): The performance metrics for each test.
             overall_avg (float): The overall average performance.
 
         Returns:
diff --git a/src/cloudai/_core/grading_strategy.py b/src/cloudai/_core/grading_strategy.py
@@ -26,8 +26,7 @@ def grade(self, directory_path: str, ideal_perf: float) -> float:
         Grades the performance of a test.
 
         Args:
-            directory_path (str): Path to the directory containing the
-                                  test's output.
+            directory_path (str): Path to the directory containing the test's output.
             ideal_perf (float): The ideal performance value for comparison.
 
         Returns:
diff --git a/src/cloudai/_core/parser.py b/src/cloudai/_core/parser.py
@@ -81,9 +81,8 @@ def parse(self) -> Tuple[System, List[TestTemplate], TestScenario]:
         Parse configurations for system, test templates, and test scenarios.
 
         Returns
-            Tuple[System, List[TestTemplate], TestScenario]: A tuple containing
-            the system object, a list of test template objects, and the test scenario
-            object.
+            Tuple[System, List[TestTemplate], TestScenario]: A tuple containing the system object, a list of test
+                template objects, and the test scenario object.
         """
         system_parser = SystemParser(self.system_config_path)
         system = system_parser.parse()
diff --git a/src/cloudai/_core/system.py b/src/cloudai/_core/system.py
@@ -22,8 +22,7 @@ class System(ABC):
 
     Attributes
         name (str): Unique name of the system.
-        scheduler (str): Type of scheduler used by the system, determining
-                         the specific subclass of System to be used.
+        scheduler (str): Type of scheduler used by the system, determining the specific subclass of System to be used.
         output_path (str): Path to the output directory.
         monitor_interval (int): Interval in seconds for monitoring jobs.
     """
@@ -54,12 +53,10 @@ def __repr__(self) -> str:
         Provide a detailed string representation of the System instance, including all its attributes.
 
         Returns
-            str: String representation of the system including name, scheduler,
-            output_path, and monitor_interval.
+            str: String representation of the system including name, scheduler, output_path, and monitor_interval.
         """
         return (
-            f"System(name='{self.name}', scheduler='{self.scheduler}', "
-            f"output_path='{self.output_path}', "
+            f"System(name='{self.name}', scheduler='{self.scheduler}', output_path='{self.output_path}', "
             f"monitor_interval={self.monitor_interval})"
         )
 
diff --git a/src/cloudai/_core/system_parser.py b/src/cloudai/_core/system_parser.py
@@ -25,8 +25,8 @@ class SystemParser:
     Parser for parsing system configurations.
 
     Attributes
-        _parsers (Dict[str, Type[BaseSystemParser]]): A mapping from system
-            types to their corresponding parser classes.
+        _parsers (Dict[str, Type[BaseSystemParser]]): A mapping from system types to their corresponding parser
+            classes.
         file_path (str): The file path to the system configuration file.
     """
 
diff --git a/src/cloudai/_core/test_parser.py b/src/cloudai/_core/test_parser.py
@@ -24,8 +24,7 @@ class TestParser(BaseMultiFileParser):
     Parser for Test objects.
 
     Attributes
-        test_template_mapping (Dict[str, TestTemplate]): Mapping of test template
-            names to TestTemplate objects.
+        test_template_mapping (Dict[str, TestTemplate]): Mapping of test template names to TestTemplate objects.
     """
 
     __test__ = False
@@ -40,8 +39,7 @@ def __init__(
 
         Args:
             directory_path (str): Path to the directory containing test data.
-            test_template_mapping (Dict[str, TestTemplate]): Mapping of test
-                template names to TestTemplate objects.
+            test_template_mapping (Dict[str, TestTemplate]): Mapping of test template names to TestTemplate objects.
         """
         super().__init__(directory_path)
         self.test_template_mapping: Dict[str, TestTemplate] = test_template_mapping
@@ -99,12 +97,11 @@ def _flatten_template_dict_keys(self, nested_args: Dict[str, Any], parent_key: s
         """
         Recursively flattens the nested dictionary structure from the test template.
 
-        Includes keys with 'default' and 'values' as valid keys, while ignoring
-        keys that specifically end with 'default' or 'values'.
+        Includes keys with 'default' and 'values' as valid keys, while ignoring keys that specifically end with
+        'default' or 'values'.
 
         Args:
-            nested_args (Dict[str, Any]): Nested argument structure from the test
-                                          template.
+            nested_args (Dict[str, Any]): Nested argument structure from the test template.
             parent_key (str): Parent key for nested arguments.
 
         Returns:
@@ -132,8 +129,7 @@ def _validate_args(self, args: Dict[str, Any], valid_keys: Set[str]) -> None:
 
         Args:
             args (Dict[str, Any]): Arguments provided in the TOML configuration.
-            valid_keys (Set[str]): Set of valid keys from the flattened template
-                                   arguments.
+            valid_keys (Set[str]): Set of valid keys from the flattened template arguments.
 
         Raises:
             ValueError: If an argument is not defined in the TestTemplate's arguments.
diff --git a/src/cloudai/_core/test_scenario.py b/src/cloudai/_core/test_scenario.py
@@ -60,8 +60,7 @@ def pretty_print(self) -> str:
                 for dep_type, dependency in test.dependencies.items():
                     if dependency:
                         s += (
-                            f"  {dep_type.replace('_', ' ').title()}: "
-                            f"{dependency.test.section_name}, "
+                            f"  {dep_type.replace('_', ' ').title()}: {dependency.test.section_name}, "
                             f"Time: {dependency.time} seconds"
                         )
             else:
diff --git a/src/cloudai/_core/test_scenario_parser.py b/src/cloudai/_core/test_scenario_parser.py
@@ -138,8 +138,7 @@ def _parse_dependencies_for_test(
         Args:
             section (str): Section name of the test.
             test_info (Dict[str, Any]): Information of the test.
-            section_tests (Dict[str, Test]): Mapping of section names to Test
-                                             objects.
+            section_tests (Dict[str, Test]): Mapping of section names to Test objects.
 
         Returns:
             Dict[str, Optional[TestDependency]]: Parsed dependencies for the test.
diff --git a/src/cloudai/_core/test_template_strategy.py b/src/cloudai/_core/test_template_strategy.py
@@ -129,8 +129,7 @@ def _override_env_vars(
 
         Args:
             default_env_vars (Dict[str, str]): The default environment variables.
-            provided_env_vars (Dict[str, str]): The provided environment variables
-                                                to override defaults.
+            provided_env_vars (Dict[str, str]): The provided environment variables to override defaults.
 
         Returns:
             Dict[str, str]: A dictionary of environment variables with overrides applied.
@@ -149,8 +148,7 @@ def _override_cmd_args(
 
         Args:
             default_cmd_args (Dict[str, str]): The default command-line arguments.
-            provided_cmd_args (Dict[str, str]): The provided command-line arguments
-                                                to override defaults.
+            provided_cmd_args (Dict[str, str]): The provided command-line arguments to override defaults.
 
         Returns:
             Dict[str, str]: A dictionary of command-line arguments with overrides applied.
diff --git a/src/cloudai/parser/system_parser/standalone_system_parser.py b/src/cloudai/parser/system_parser/standalone_system_parser.py
@@ -33,9 +33,8 @@ def parse(self, data: Dict[str, Any]) -> StandaloneSystem:
             StandaloneSystem: The parsed Standalone system object.
 
         Raises:
-            ValueError: If 'name' or 'output_path' are missing from
-            the data or if there are node list parsing issues or group
-            membership conflicts.
+            ValueError: If 'name' or 'output_path' are missing from the data or if there are node list parsing issues
+            or group membership conflicts.
         """
         name = data.get("name")
         if not name:
diff --git a/src/cloudai/report_generator/report_generator.py b/src/cloudai/report_generator/report_generator.py
@@ -23,8 +23,8 @@ class ReportGenerator:
     """
     Generates reports for each test in a TestScenario.
 
-    By identifying the appropriate directories for each test and using test templates to
-    generate detailed reports based on subdirectories.
+    By identifying the appropriate directories for each test and using test templates to generate detailed reports
+    based on subdirectories.
     """
 
     def __init__(self, output_path: str) -> None:
@@ -40,8 +40,8 @@ def generate_report(self, test_scenario: TestScenario) -> None:
         """
         Iterate over tests in the given test scenario.
 
-        Identifies the relevant directories based on the test's section name, and generates
-        a report for each test using its associated test template.
+        Identifies the relevant directories based on the test's section name, and generates a report for each test
+        using its associated test template.
 
         Args:
             test_scenario (TestScenario): The scenario containing tests.
diff --git a/src/cloudai/runner/slurm/slurm_runner.py b/src/cloudai/runner/slurm/slurm_runner.py
@@ -31,17 +31,13 @@ class SlurmRunner(BaseRunner):
     """
     Implementation of the Runner for a system using Slurm.
 
-    This class is responsible for executing and managing tests in a Slurm
-    environment. It extends the BaseRunner class, implementing the abstract
-    methods to work with Slurm jobs.
+    This class is responsible for executing and managing tests in a Slurm environment. It extends the BaseRunner class,
+    implementing the abstract methods to work with Slurm jobs.
 
     Attributes
-        slurm_system (SlurmSystem): This attribute is a casted version of the `system`
-                                    attribute to `SlurmSystem` type, ensuring that
-                                    Slurm-specific properties and methods are
-                                    accessible.
-        cmd_shell (CommandShell): An instance of CommandShell for executing
-                                  system commands.
+        slurm_system (SlurmSystem): This attribute is a casted version of the `system` attribute to `SlurmSystem` type,
+            ensuring that Slurm-specific properties and methods are accessible.
+        cmd_shell (CommandShell): An instance of CommandShell for executing system commands.
         Inherits all other attributes from the BaseRunner class.
     """
 
diff --git a/src/cloudai/runner/standalone/standalone_runner.py b/src/cloudai/runner/standalone/standalone_runner.py
@@ -30,13 +30,11 @@ class StandaloneRunner(BaseRunner):
     """
     Implementation of the Runner for a system using Standalone.
 
-    This class is responsible for executing and managing tests in a standalone
-    environment. It extends the BaseRunner class, implementing the abstract
-    methods to work with standalone jobs.
+    This class is responsible for executing and managing tests in a standalone environment. It extends the BaseRunner
+    class, implementing the abstract methods to work with standalone jobs.
 
     Attributes
-        cmd_shell (CommandShell): An instance of CommandShell for executing
-                                  system commands.
+        cmd_shell (CommandShell): An instance of CommandShell for executing system commands.
         Inherits all other attributes from the BaseRunner class.
     """
 
diff --git a/src/cloudai/schema/test_template/common/slurm_job_id_retrieval_strategy.py b/src/cloudai/schema/test_template/common/slurm_job_id_retrieval_strategy.py
@@ -22,16 +22,14 @@ class SlurmJobIdRetrievalStrategy(JobIdRetrievalStrategy):
     """
     Strategy for retrieving job IDs from Slurm job submission outputs.
 
-    Implements JobIdRetrievalStrategy to extract job IDs from the standard
-    output of Slurm scheduler submissions.
+    Implements JobIdRetrievalStrategy to extract job IDs from the standard output of Slurm scheduler submissions.
     """
 
     def get_job_id(self, stdout: str, stderr: str) -> Optional[int]:
         """
         Extract the job ID from Slurm command output.
 
-        Parses stdout for a pattern matching Slurm submission confirmation
-        and extracts the job ID.
+        Parses stdout for a pattern matching Slurm submission confirmation and extracts the job ID.
 
         Args:
             stdout (str): Standard output from the Slurm command.
diff --git a/src/cloudai/schema/test_template/common/standalone_job_id_retrieval_strategy.py b/src/cloudai/schema/test_template/common/standalone_job_id_retrieval_strategy.py
@@ -21,9 +21,8 @@ class StandaloneJobIdRetrievalStrategy(JobIdRetrievalStrategy):
     """
     Strategy for retrieving job IDs from standalone application outputs.
 
-    This class provides a specific implementation of JobIdRetrievalStrategy
-    for extracting job IDs from the output of standalone applications, assuming
-    the job ID is directly available in the standard output.
+    This class provides a specific implementation of JobIdRetrievalStrategy for extracting job IDs from the output of
+    standalone applications, assuming the job ID is directly available in the standard output.
     """
 
     def get_job_id(self, stdout: str, stderr: str) -> Optional[int]:
@@ -37,8 +36,7 @@ def get_job_id(self, stdout: str, stderr: str) -> Optional[int]:
             stderr (str): The standard error from the application (unused).
 
         Returns:
-            Optional[int]: The extracted job ID as an integer, or None if
-            the conversion fails due to invalid output.
+            Optional[int]: The extracted job ID as an integer, or None if the conversion fails due to invalid output.
         """
         try:
             return int(stdout)
diff --git a/src/cloudai/schema/test_template/jax_toolbox/slurm_command_gen_strategy.py b/src/cloudai/schema/test_template/jax_toolbox/slurm_command_gen_strategy.py
@@ -63,8 +63,7 @@ def _format_xla_flags(self, cmd_args: Dict[str, str]) -> str:
             cmd_args (Dict[str, str]): Command-line arguments.
 
         Returns:
-            str: A single string containing all XLA-related flags formatted for
-                 inclusion in the environment variables.
+            str: A single string containing all XLA-related flags formatted for inclusion in the environment variables.
         """
         xla_flags = []
 
@@ -153,12 +152,11 @@ def _create_run_script(
         """
         Generate and writes the run.sh script to the specified output directory.
 
-        The script configures environment variables, applies necessary command
-        options, and executes the Python command within the SLURM environment.
+        The script configures environment variables, applies necessary command options, and executes the Python command
+        within the SLURM environment.
 
         Args:
-            slurm_args (Dict[str, Any]): SLURM arguments including the output path
-                                         and other job-related settings.
+            slurm_args (Dict[str, Any]): SLURM arguments including the output path and other job-related settings.
             env_vars (Dict[str, str]): Environment variables.
             cmd_args (Dict[str, str]): Command-line arguments.
             extra_cmd_args (str): Additional command-line arguments to be included
@@ -249,12 +247,10 @@ def _generate_python_command(
 
         Args:
             stage (str): The stage of processing (e.g., 'profile', 'perf').
-            slurm_args (Dict[str, Any]): Dictionary containing the SLURM job settings
-                                         such as number of nodes.
+            slurm_args (Dict[str, Any]): Dictionary containing the SLURM job settings such as number of nodes.
             env_vars (Dict[str, str]): Environment variables.
             cmd_args (Dict[str, str]): Command-line arguments.
-            extra_cmd_args (str): Additional command-line arguments to be included
-                                  in the Python command.
+            extra_cmd_args (str): Additional command-line arguments to be included in the Python command.
 
         Returns:
             str: The formatted Python command string to be executed within a
diff --git a/src/cloudai/schema/test_template/nccl_test/grading_strategy.py b/src/cloudai/schema/test_template/nccl_test/grading_strategy.py
@@ -21,9 +21,8 @@ class NcclTestGradingStrategy(GradingStrategy):
     """
     Performance grading strategy for NcclTest test templates on Slurm systems.
 
-    Evaluates the test's performance by comparing the maximum bus bandwidth
-    achieved during the test against the test's ideal performance metric.
-    The grade is normalized and scaled between 0 and 100.
+    Evaluates the test's performance by comparing the maximum bus bandwidth achieved during the test against the test's
+    ideal performance metric. The grade is normalized and scaled between 0 and 100.
     """
 
     def grade(self, directory_path: str, ideal_perf: float) -> float:
@@ -58,8 +57,7 @@ def _extract_max_bus_bandwidth(self, stdout_path: str) -> float:
         Extract the maximum bus bandwidth from the NcclTest output file.
 
         Args:
-            stdout_path (str): Path to the stdout.txt file containing the
-                               NcclTest output.
+            stdout_path (str): Path to the stdout.txt file containing the NcclTest output.
 
         Returns:
             float: The maximum bus bandwidth value.
diff --git a/src/cloudai/schema/test_template/nemo_launcher/grading_strategy.py b/src/cloudai/schema/test_template/nemo_launcher/grading_strategy.py
@@ -25,8 +25,7 @@ def grade(self, directory_path: str, ideal_perf: float) -> float:
         Grades the performance of a test.
 
         Args:
-            directory_path (str): Path to the directory containing the
-                                  test's output.
+            directory_path (str): Path to the directory containing the test's output.
             ideal_perf (float): The ideal performance value for comparison.
 
         Returns:
diff --git a/src/cloudai/schema/test_template/nemo_launcher/slurm_job_id_retrieval_strategy.py b/src/cloudai/schema/test_template/nemo_launcher/slurm_job_id_retrieval_strategy.py
@@ -20,20 +20,18 @@
 
 class NeMoLauncherSlurmJobIdRetrievalStrategy(JobIdRetrievalStrategy):
     """
-    Strategy for retrieving job IDs from NeMo Megatron launcher submissions to a Slurm scheduler.
+    Strategy for retrieving job IDs from NeMo launcher submissions to a Slurm scheduler.
 
-    This class implements the JobIdRetrievalStrategy interface to extract job IDs
-    from the standard output provided by the Slurm scheduler when the NeMo Megatron
-    launcher template is used for job submission.
+    This class implements the JobIdRetrievalStrategy interface to extract job IDs from the standard output provided by
+    the Slurm scheduler when the NeMo launcher template is used for job submission.
     """
 
     def get_job_id(self, stdout: str, stderr: str) -> Optional[int]:
         """
         Extract the job ID from the Slurm command output.
 
-        This method searches the standard output for a specific pattern that
-        matches the submission confirmation message of a Slurm job and extracts
-        the job ID.
+        This method searches the standard output for a specific pattern that matches the submission confirmation
+        message of a Slurm job and extracts the job ID.
 
         Args:
             stdout (str): The standard output from the Slurm command.
diff --git a/src/cloudai/schema/test_template/nemo_launcher/template.py b/src/cloudai/schema/test_template/nemo_launcher/template.py
diff --git a/src/cloudai/schema/test_template/sleep/grading_strategy.py b/src/cloudai/schema/test_template/sleep/grading_strategy.py
diff --git a/src/cloudai/schema/test_template/sleep/standalone_command_gen_strategy.py b/src/cloudai/schema/test_template/sleep/standalone_command_gen_strategy.py
diff --git a/src/cloudai/schema/test_template/sleep/standalone_install_strategy.py b/src/cloudai/schema/test_template/sleep/standalone_install_strategy.py
diff --git a/src/cloudai/schema/test_template/ucc_test/grading_strategy.py b/src/cloudai/schema/test_template/ucc_test/grading_strategy.py
diff --git a/src/cloudai/systems/slurm/slurm_node.py b/src/cloudai/systems/slurm/slurm_node.py
diff --git a/src/cloudai/systems/slurm/slurm_system.py b/src/cloudai/systems/slurm/slurm_system.py
diff --git a/src/cloudai/systems/slurm/strategy/slurm_command_gen_strategy.py b/src/cloudai/systems/slurm/strategy/slurm_command_gen_strategy.py
diff --git a/src/cloudai/util/command_shell.py b/src/cloudai/util/command_shell.py
