Merge pull request #732 from NVIDIA/am/system-inheritance

Convert base System into pydantic model

Author: amaslenn

src/cloudai/_core/system.py

@@ -18,67 +18,27 @@
 import logging
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, Optional
+from typing import TYPE_CHECKING, Any
+
+from pydantic import BaseModel, ConfigDict, Field
 
 from .installables import Installable
 
 if TYPE_CHECKING:
     from .base_job import BaseJob
 
 
-class System(ABC):
-    """
-    Base class representing a generic system.
-
-    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.
-        install_path (Path): Installation path of CloudAI software.
-        output_path (Path): Path to the output directory.
-        global_env_vars (Optional[Dict[str, Any]]): Dictionary containing additional configuration settings for the
-            system.
-        monitor_interval (int): Interval in seconds for monitoring jobs.
-    """
-
-    def __init__(
-        self,
-        name: str,
-        scheduler: str,
-        install_path: Path,
-        output_path: Path,
-        global_env_vars: Optional[Dict[str, Any]] = None,
-        monitor_interval: int = 1,
-    ) -> None:
-        """
-        Initialize a System instance.
+class System(ABC, BaseModel):
+    """Base class representing a generic system."""
 
-        Args:
-            name (str): Name of the system.
-            scheduler (str): Type of scheduler used by the system.
-            install_path (Path): The installation path of CloudAI.
-            output_path (Path): Path to the output directory.
-            global_env_vars (Optional[Dict[str, Any]]): Dictionary containing additional configuration settings for
-                the system.
-            monitor_interval (int): Interval in seconds for monitoring jobs.
-        """
-        self.name = name
-        self.scheduler = scheduler
-        self.install_path = install_path
-        self.output_path = output_path
-        self.global_env_vars = global_env_vars if global_env_vars is not None else {}
-        self.monitor_interval = monitor_interval
-
-    def __repr__(self) -> str:
-        """
-        Provide a detailed string representation of the System instance, including all its attributes.
+    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
 
-        Returns
-            str: String representation of the system including name, scheduler, output_path, and monitor_interval.
-        """
-        return (
-            f"System(name='{self.name}', scheduler='{self.scheduler}', output_path='{self.output_path}', "
-            f"monitor_interval={self.monitor_interval})"
-        )
+    name: str
+    scheduler: str
+    install_path: Path
+    output_path: Path
+    global_env_vars: dict[str, Any] = Field(default_factory=dict)
+    monitor_interval: int = 1
 
     @abstractmethod
     def update(self) -> None:

src/cloudai/systems/kubernetes/kubernetes_system.py

@@ -27,41 +27,18 @@
     import kubernetes as k8s
 
 
-from pydantic import BaseModel, ConfigDict
-
 from cloudai.core import BaseJob, System
 from cloudai.util.lazy_imports import lazy
 
 from .kubernetes_job import KubernetesJob
 
 
-class KubernetesSystem(BaseModel, System):
-    """
-    Represents a Kubernetes system.
-
-    Attributes
-        name (str): The name of the Kubernetes system.
-        install_path (Path): Path to the installation directory.
-        output_path (Path): Path to the output directory.
-        kube_config_path (Path): Path to the Kubernetes config file.
-        default_namespace (str): The default Kubernetes namespace for jobs.
-        scheduler (str): The scheduler type, default is "kubernetes".
-        global_env_vars (Dict[str, Any]): Global environment variables to be passed to jobs.
-        monitor_interval (int): Time interval to monitor jobs, in seconds.
-        _core_v1 (client.CoreV1Api): Kubernetes Core V1 API client instance.
-        _batch_v1 (client.BatchV1Api): Kubernetes Batch V1 API client instance.
-        _custom_objects_api (CustomObjectsApi): Kubernetes Custom Objects API client instance.
-    """
-
-    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
-
-    name: str
-    install_path: Path
-    output_path: Path
+class KubernetesSystem(System):
+    """Represents a Kubernetes system."""
+
     kube_config_path: Path
     default_namespace: str
     scheduler: str = "kubernetes"
-    global_env_vars: Dict[str, Any] = {}
     monitor_interval: int = 1
     _core_v1: Optional[k8s.client.CoreV1Api] = None
     _batch_v1: Optional[k8s.client.BatchV1Api] = None

src/cloudai/systems/lsf/lsf_system.py

@@ -16,7 +16,7 @@
 
 import logging
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
 from pydantic import BaseModel, ConfigDict, Field, field_serializer
 
@@ -49,7 +49,7 @@ class LSFQueue(BaseModel):
     lsf_nodes: List[LSFNodeObj] = Field(default_factory=list, exclude=True)
 
 
-class LSFSystem(BaseModel, System):
+class LSFSystem(System):
     """
     Represents an LSF system.
 
@@ -69,14 +69,8 @@ class LSFSystem(BaseModel, System):
         cmd_shell (CommandShell): Command shell for executing system commands.
     """
 
-    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
-
-    name: str
-    install_path: Path
-    output_path: Path
     queues: List[LSFQueue] = Field(default_factory=list)
     account: Optional[str] = None
-    global_env_vars: Dict[str, Any] = {}
     scheduler: str = "lsf"
     project_name: Optional[str] = None
     default_queue: Optional[str] = None

src/cloudai/systems/runai/runai_system.py

@@ -18,7 +18,7 @@
 from pathlib import Path
 from typing import Any, Dict, List, Optional
 
-from pydantic import BaseModel, ConfigDict, PrivateAttr
+from pydantic import Field, PrivateAttr
 
 from cloudai.core import BaseJob, System
 
@@ -31,7 +31,7 @@
 from .runai_training import ActualPhase, RunAITraining
 
 
-class RunAISystem(BaseModel, System):
+class RunAISystem(System):
     """
     RunAISystem integrates with the RunAI platform to manage and monitor jobs and nodes.
 
@@ -47,21 +47,15 @@ class RunAISystem(BaseModel, System):
         nodes (List[RunAINode]): List of nodes in the RunAI cluster.
     """
 
-    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
-
-    name: str
     scheduler: str = "runai"
-    install_path: Path
-    output_path: Path
-    global_env_vars: Dict[str, Any] = {}
     monitor_interval: int = 60
     base_url: str
     user_email: str
     app_id: str
     app_secret: str
     project_id: str
     cluster_id: str
-    nodes: List[RunAINode] = []
+    nodes: List[RunAINode] = Field(default_factory=list)
     _api_client: Optional[RunAIRestClient] = PrivateAttr(default=None)
 
     @property

src/cloudai/systems/slurm/slurm_system.py

@@ -93,32 +93,9 @@ class SlurmPartition(BaseModel):
     slurm_nodes: list[SlurmNode] = Field(default_factory=list[SlurmNode], exclude=True)
 
 
-class SlurmSystem(BaseModel, System):
-    """
-    Represents a Slurm system.
-
-    Attributes
-        output_path (Path): Path to the output directory.
-        default_partition (str): The default partition for job submission.
-        partitions (Dict[str, List[SlurmNode]]): Mapping of partition names to lists of SlurmNodes.
-        account (Optional[str]): Account name for charging resources used by this job.
-        distribution (Optional[str]): Specifies alternate distribution methods for remote processes.
-        mpi (Optional[str]): Indicates the Process Management Interface (PMI) implementation to be used for
-            inter-process communication.
-        gpus_per_node (Optional[int]): Specifies the number of GPUs available per node.
-        ntasks_per_node (Optional[int]): Specifies the number of tasks that can run concurrently on a single node.
-        cache_docker_images_locally (bool): Whether to cache Docker images locally for the Slurm system.
-        groups (Dict[str, Dict[str, List[SlurmNode]]]): Nested mapping where the key is the partition name and the
-            value is another dictionary with group names as keys and lists of SlurmNodes as values, representing the
-            group composition within each partition.
-        cmd_shell (CommandShell): An instance of CommandShell for executing system commands.
-    """
+class SlurmSystem(System):
+    """Represents a Slurm system."""
 
-    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
-
-    name: str
-    install_path: Path
-    output_path: Path
     default_partition: str
     partitions: List[SlurmPartition]
     account: Optional[str] = None
@@ -127,12 +104,11 @@ class SlurmSystem(BaseModel, System):
     gpus_per_node: Optional[int] = None
     ntasks_per_node: Optional[int] = None
     cache_docker_images_locally: bool = False
-    global_env_vars: Dict[str, Any] = {}
     scheduler: str = "slurm"
     monitor_interval: int = 60
     cmd_shell: CommandShell = Field(default=CommandShell(), exclude=True)
     extra_srun_args: Optional[str] = None
-    extra_sbatch_args: list[str] = []
+    extra_sbatch_args: list[str] = Field(default_factory=list)
     supports_gpu_directives_cache: Optional[bool] = Field(default=None, exclude=True)
     container_mount_home: bool = False
 

src/cloudai/systems/standalone/standalone_system.py

@@ -15,26 +15,18 @@
 # limitations under the License.
 
 import logging
-from pathlib import Path
-
-from pydantic import BaseModel, ConfigDict
 
 from cloudai.core import BaseJob, System
 from cloudai.util import CommandShell
 
 
-class StandaloneSystem(BaseModel, System):
+class StandaloneSystem(System):
     """
     Class representing a Standalone system.
 
     This class is used for systems that execute commands directly without a job scheduler.
     """
 
-    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
-
-    name: str
-    install_path: Path
-    output_path: Path
     scheduler: str = "standalone"
     monitor_interval: int = 1
     cmd_shell: CommandShell = CommandShell()