FedML-AI Docstrings Update

python/examples/launch/serve_mnist/model/minist_model.py

@@ -1,11 +1,54 @@
 import torch
 class LogisticRegression(torch.nn.Module):
+    """
+    Logistic Regression model for binary classification.
+
+    This class defines a logistic regression model with a single linear layer followed by a sigmoid activation function
+    for binary classification tasks.
+
+    Args:
+        input_dim (int): The dimensionality of the input features.
+        output_dim (int): The number of output classes, which should be 1 for binary classification.
+
+    Example:
+        # Create a logistic regression model for binary classification
+        input_dim = 10
+        output_dim = 1
+        model = LogisticRegression(input_dim, output_dim)
+
+    Forward Method:
+        The forward method computes the output of the model for a given input.
+
+    Example:
+        # Forward pass with input tensor 'x'
+        input_tensor = torch.tensor([0.1, 0.2, 0.3, ..., 0.9])
+        output = model(input_tensor)
+
+    Note:
+        - For binary classification, the `output_dim` should be set to 1.
+        - The `forward` method applies a sigmoid activation to the linear output, producing values in the range [0, 1].
+
+    """
+
     def __init__(self, input_dim, output_dim):
         super(LogisticRegression, self).__init__()
         self.linear = torch.nn.Linear(input_dim, output_dim)
 
     def forward(self, x):
-        import torch
+        """
+        Forward pass of the logistic regression model.
+
+        Args:
+            x (torch.Tensor): Input tensor of shape (batch_size, input_dim).
+
+        Returns:
+            torch.Tensor: Output tensor of shape (batch_size, output_dim).
+
+        Example:
+            # Forward pass with input tensor 'x'
+            input_tensor = torch.tensor([0.1, 0.2, 0.3, ..., 0.9])
+            output = model(input_tensor)
+        """
         outputs = torch.sigmoid(self.linear(x))
         return outputs
 

python/fedml/__init__.py

@@ -22,6 +22,56 @@
 )
 from .core.common.ml_engine_backend import MLEngineBackend
 
+from fedml import device
+from fedml import data
+from fedml import model
+from fedml import mlops
+
+from .arguments import load_arguments
+
+from .launch_simulation import run_simulation
+
+from .launch_cross_silo_horizontal import run_cross_silo_server
+from .launch_cross_silo_horizontal import run_cross_silo_client
+
+from .launch_cross_silo_hi import run_hierarchical_cross_silo_server
+from .launch_cross_silo_hi import run_hierarchical_cross_silo_client
+
+from .launch_cheeath import run_cheetah_server
+from .launch_cheeath import run_cheetah_client
+
+from .launch_serving import run_model_serving_client
+from .launch_serving import run_model_serving_server
+
+from .launch_cross_device import run_mnn_server
+
+from .core.common.ml_engine_backend import MLEngineBackend
+
+from .runner import FedMLRunner
+
+from fedml import api
+
+__all__ = [
+    "MLEngineBackend",
+    "device",
+    "data",
+    "model",
+    "mlops",
+    "FedMLRunner",
+    "run_simulation",
+    "run_cross_silo_server",
+    "run_cross_silo_client",
+    "run_hierarchical_cross_silo_server",
+    "run_hierarchical_cross_silo_client",
+    "run_cheetah_server",
+    "run_cheetah_client",
+    "run_model_serving_client",
+    "run_model_serving_server",
+    "run_mnn_server",
+    "api"
+]
+
+
 _global_training_type = None
 _global_comm_backend = None
 
@@ -32,7 +82,17 @@ def init(args=None, check_env=True, should_init_logs=True):
     if args is None:
         args = load_arguments(fedml._global_training_type, fedml._global_comm_backend)
 
-    """Initialize FedML Engine."""
+    """
+    Initialize the FedML Engine.
+
+    Args:
+        args (argparse.Namespace, optional): Command-line arguments. Defaults to None.
+        check_env (bool, optional): Whether to check the environment. Defaults to True.
+        should_init_logs (bool, optional): Whether to initialize logs. Defaults to True.
+
+    Returns:
+        argparse.Namespace: Updated command-line arguments.
+    """
     if check_env:
         collect_env(args)
 
@@ -120,6 +180,12 @@ def init(args=None, check_env=True, should_init_logs=True):
 
 
 def print_args(args):
+    """
+    Print the arguments to the log, excluding sensitive paths.
+
+    Args:
+        args (argparse.Namespace): Command-line arguments.
+    """
     mqtt_config_path = None
     s3_config_path = None
     args_copy = args
@@ -138,7 +204,9 @@ def print_args(args):
 
 def update_client_specific_args(args):
     """
-        data_silo_config is used for reading specific configuration for each client
+    Update client-specific arguments based on data_silo_config.
+
+    data_silo_config is used for reading specific configuration for each client
         Example: In fedml_config.yaml, we have the following configuration
         client_specific_args: 
             data_silo_config: 
@@ -149,6 +217,9 @@ def update_client_specific_args(args):
                     fedml_config/data_silo_4_config.yaml,
                 ]
             data_silo_1_config.yaml contains some client client speicifc arguments.
+
+    Args:
+        args (argparse.Namespace): Command-line arguments.
     """
     if (
         hasattr(args, "data_silo_config")
@@ -166,7 +237,17 @@ def update_client_specific_args(args):
 
 
 def init_simulation_mpi(args):
+
     from mpi4py import MPI
+    """
+    Initialize MPI-based simulation.
+
+    Args:
+        args (argparse.Namespace): Command-line arguments.
+
+    Returns:
+        argparse.Namespace: Updated command-line arguments.
+    """
 
     comm = MPI.COMM_WORLD
     process_id = comm.Get_rank()
@@ -183,14 +264,35 @@ def init_simulation_mpi(args):
 
 
 def init_simulation_sp(args):
+    """
+    Initialize single-process simulation.
+
+    Args:
+        args (argparse.Namespace): Command-line arguments.
+
+    Returns:
+        argparse.Namespace: Updated command-line arguments.
+    """
     return args
 
 
 def init_simulation_nccl(args):
+    """
+    Initialize NCCL-based simulation.
+
+    Args:
+        args (argparse.Namespace): Command-line arguments.
+    """
     return
 
 
 def manage_profiling_args(args):
+    """
+    Manage profiling-related arguments and configurations.
+
+    Args:
+        args (argparse.Namespace): Command-line arguments.
+    """
     if not hasattr(args, "sys_perf_profiling"):
         args.sys_perf_profiling = True
     if not hasattr(args, "sys_perf_profiling"):
@@ -236,6 +338,12 @@ def manage_profiling_args(args):
 
 
 def manage_cuda_rpc_args(args):
+    """
+    Manage CUDA RPC-related arguments and configurations.
+
+    Args:
+        args (argparse.Namespace): Command-line arguments.
+    """
 
     if (not hasattr(args, "enable_cuda_rpc")) or (not args.using_gpu):
         args.enable_cuda_rpc = False
@@ -264,6 +372,12 @@ def manage_cuda_rpc_args(args):
 
 
 def manage_mpi_args(args):
+    """
+    Manage MPI-related arguments and configurations.
+
+    Args:
+        args (argparse.Namespace): Command-line arguments.
+    """
     if hasattr(args, "backend") and args.backend == "MPI":
         from mpi4py import MPI
 
@@ -282,6 +396,15 @@ def manage_mpi_args(args):
         args.comm = None
 
 def init_cross_silo_horizontal(args):
+    """
+    Initialize the cross-silo training for the horizontal scenario.
+
+    Args:
+        args (argparse.Namespace): Command-line arguments.
+
+    Returns:
+        args (argparse.Namespace): Updated command-line arguments.
+    """
     args.n_proc_in_silo = 1
     args.proc_rank_in_silo = 0
     manage_mpi_args(args)
@@ -291,6 +414,15 @@ def init_cross_silo_horizontal(args):
 
 
 def init_cross_silo_hierarchical(args):
+    """
+    Initialize the cross-silo training for the hierarchical scenario.
+
+    Args:
+        args (argparse.Namespace): Command-line arguments.
+
+    Returns:
+        args (argparse.Namespace): Updated command-line arguments.
+    """
     manage_mpi_args(args)
     manage_cuda_rpc_args(args)
 
@@ -344,6 +476,15 @@ def init_cross_silo_hierarchical(args):
 
 
 def init_cheetah(args):
+    """
+    Initialize the CheetaH training scenario.
+
+    Args:
+        args (argparse.Namespace): Command-line arguments.
+
+    Returns:
+        args (argparse.Namespace): Updated command-line arguments.
+    """
     args.n_proc_in_silo = 1
     args.proc_rank_in_silo = 0
     manage_mpi_args(args)
@@ -353,6 +494,15 @@ def init_cheetah(args):
 
 
 def init_model_serving(args):
+    """
+    Initialize the model serving scenario.
+
+    Args:
+        args (argparse.Namespace): Command-line arguments.
+
+    Returns:
+        args (argparse.Namespace): Updated command-line arguments.
+    """
     args.n_proc_in_silo = 1
     args.proc_rank_in_silo = 0
     manage_cuda_rpc_args(args)
@@ -361,10 +511,12 @@ def init_model_serving(args):
 
 
 def update_client_id_list(args):
-
     """
-        generate args.client_id_list for CLI mode where args.client_id_list is set to None
-        In MLOps mode, args.client_id_list will be set to real-time client id list selected by UI (not starting from 1)
+    Generate args.client_id_list for the CLI mode where args.client_id_list is set to None.
+    In MLOps mode, args.client_id_list will be set to a real-time client id list selected by the UI (not starting from 1).
+
+    Args:
+        args (argparse.Namespace): Command-line arguments.
     """
     if not hasattr(args, "using_mlops") or (hasattr(args, "using_mlops") and not args.using_mlops):
         if not hasattr(args, "client_id_list") or args.client_id_list is None or args.client_id_list == "None" or args.client_id_list == "[]":
@@ -396,60 +548,22 @@ def update_client_id_list(args):
 
 
 def init_cross_device(args):
+    """
+    Initialize the cross-device training scenario.
+
+    Args:
+        args (argparse.Namespace): Command-line arguments.
+
+    Returns:
+        args (argparse.Namespace): Updated command-line arguments.
+    """
     args.rank = 0  # only server runs on Python package
     args.role = "server"
     return args
 
 
 def run_distributed():
+    """
+    Placeholder function for running distributed training.
+    """
     pass
-
-
-from fedml import device
-from fedml import data
-from fedml import model
-from fedml import mlops
-
-from .arguments import load_arguments
-
-from .launch_simulation import run_simulation
-
-from .launch_cross_silo_horizontal import run_cross_silo_server
-from .launch_cross_silo_horizontal import run_cross_silo_client
-
-from .launch_cross_silo_hi import run_hierarchical_cross_silo_server
-from .launch_cross_silo_hi import run_hierarchical_cross_silo_client
-
-from .launch_cheeath import run_cheetah_server
-from .launch_cheeath import run_cheetah_client
-
-from .launch_serving import run_model_serving_client
-from .launch_serving import run_model_serving_server
-
-from .launch_cross_device import run_mnn_server
-
-from .core.common.ml_engine_backend import MLEngineBackend
-
-from .runner import FedMLRunner
-
-from fedml import api
-
-__all__ = [
-    "MLEngineBackend",
-    "device",
-    "data",
-    "model",
-    "mlops",
-    "FedMLRunner",
-    "run_simulation",
-    "run_cross_silo_server",
-    "run_cross_silo_client",
-    "run_hierarchical_cross_silo_server",
-    "run_hierarchical_cross_silo_client",
-    "run_cheetah_server",
-    "run_cheetah_client",
-    "run_model_serving_client",
-    "run_model_serving_server",
-    "run_mnn_server",
-    "api"
-]