Merge pull request #159 from DanielAvdar/docs

Refactor documentation and improve code clarity across modules.

Author: DanielAvdar

Makefile

@@ -12,6 +12,9 @@ test:
 coverage:
 	poetry run pytest --cov=ml_orchestrator --cov-report=xml --junitxml=junit.xml -o junit_family=legacy
 
+cov:
+	poetry run pytest --cov=ml_orchestrator --cov-report=term-missing
+
 check:
 	poetry run pre-commit run --all-files
 mypy:

docs/source/usage.rst

@@ -7,13 +7,14 @@ This guide provides a comprehensive walkthrough for utilizing the **ml-orchestra
 
 Installation
 ------------
-To install the package, use `Poetry` with the following command:
+
 
 .. code-block:: bash
 
-   pip install ml-orchestrator
+   pip install ml-orchestrator[editor]
 
-The **ml-orchestrator** package has no external dependencies by default.
+.. note::
+   The core **ml-orchestrator** package is intentionally designed to be dependency-free. This architectural decision ensures that components created with ml-orchestrator won't have unnecessary liabilities or dependencies, making them more portable and easier to maintain in production environments.
 
 Quick Start
 -----------

ml_orchestrator/__init__.py

@@ -1,3 +1,9 @@
+"""ML Orchestrator package for machine learning pipeline components management.
+
+This package provides tools for creating, parsing, and managing machine learning
+pipeline components with a focus on Kubeflow Pipelines integration.
+"""
+
 from importlib.metadata import version
 
 from .comp_parser import ComponentParser

ml_orchestrator/artifacts/__init__.py

@@ -1,3 +1,10 @@
+"""Artifacts package for ML pipeline component inputs and outputs.
+
+This package provides classes and types for representing and managing
+different kinds of machine learning artifacts such as models, datasets,
+and metrics, with support for JSON serialization.
+"""
+
 from typing import Annotated, TypeVar
 
 from ml_orchestrator.artifacts.artifact import Artifact

ml_orchestrator/artifacts/artifact.py

@@ -1,9 +1,32 @@
+"""Module for defining base artifact classes used in ML pipeline components.
+
+This module provides base classes for representing and managing artifacts
+with support for JSON serialization of metadata.
+"""
+
 import json
 from typing import Any, Dict, Mapping, Optional
 
 
 class JSONSerializableDict(dict):
+    """Dictionary that ensures all values are JSON serializable.
+
+    A specialized dictionary that validates all values to ensure they can be
+    serialized to JSON format. This is useful for storing metadata that
+    needs to be serialized for storage or transmission.
+    """
+
     def __setitem__(self, key: Any, value: Any) -> None:
+        """Set an item in the dictionary after ensuring it's JSON serializable.
+
+        Args:
+            key: The dictionary key
+            value: The value to set, must be JSON serializable
+
+        Raises:
+            ValueError: If the value cannot be serialized to JSON
+
+        """
         try:
             json.dumps(value)
             return super().__setitem__(key, value)
@@ -13,6 +36,16 @@ def __setitem__(self, key: Any, value: Any) -> None:
         raise ValueError(f"Value for key '{key}' is not JSON serializable: {value}")
 
     def update(self, m: Mapping = None, **kwargs: Any) -> None:  # type: ignore[override]
+        """Update the dictionary with another mapping after ensuring values are JSON serializable.
+
+        Args:
+            m: A mapping to update from
+            **kwargs: Additional key-value pairs to update with
+
+        Raises:
+            ValueError: If any values cannot be serialized to JSON
+
+        """
         try:
             json.dumps(m)
             return super().update(m, **kwargs)
@@ -22,30 +55,79 @@ def update(self, m: Mapping = None, **kwargs: Any) -> None:  # type: ignore[over
         raise ValueError(f"Aritfact metadata must be JSON serializable, got: {m}")
 
     def __or__(self, other: Mapping) -> "JSONSerializableDict":
+        """Implement the | operator to combine two mappings.
+
+        Args:
+            other: Another mapping to combine with this one
+
+        Returns:
+            A new JSONSerializableDict containing items from both mappings
+
+        """
         new_dict = JSONSerializableDict(self)
         new_dict.update(other)
         return new_dict
 
 
 class Artifact:
+    """Base class for all artifacts used in ML pipeline components.
+
+    An artifact represents a file or directory that is an input to or output from
+    a pipeline component, with associated metadata stored in a JSON-serializable format.
+
+    Attributes:
+        schema_title: The schema title for this artifact type
+        schema_version: The schema version for this artifact type
+
+    """
+
     schema_title = "system.Artifact"
     schema_version = "0.0.1"
 
     def __init__(self, name: Optional[str] = None, uri: Optional[str] = None, metadata: Optional[Dict] = None) -> None:
+        """Initialize an artifact with a name, URI, and metadata.
+
+        Args:
+            name: The name of the artifact
+            uri: The URI where the artifact is stored
+            metadata: A dictionary of metadata to associate with the artifact
+
+        """
         self.uri = uri or "./"
         self.name = name or "artifact"
         self._metadata = JSONSerializableDict(metadata or dict())
 
     @property
     def path(self) -> str:
+        """Get the path where the artifact is stored.
+
+        Returns:
+            The URI of the artifact
+
+        """
         return self.uri
 
     @property
     def metadata(self) -> "JSONSerializableDict":
+        """Get the metadata associated with this artifact.
+
+        Returns:
+            A JSON-serializable dictionary of metadata
+
+        """
         return self._metadata
 
     @metadata.setter
     def metadata(self, new_data: Mapping) -> None:
+        """Set new metadata for this artifact.
+
+        Args:
+            new_data: A mapping containing the new metadata
+
+        Raises:
+            ValueError: If the new metadata cannot be serialized to JSON
+
+        """
         try:
             json.dumps(new_data)
             self._metadata = JSONSerializableDict(new_data)

ml_orchestrator/artifacts/artifacts.py

@@ -1,3 +1,9 @@
+"""Module for specialized artifact types used in ML pipeline components.
+
+This module defines various specialized artifact types that extend the base Artifact
+class, providing specific implementations for different kinds of ML outputs and inputs.
+"""
+
 from typing import Dict, Optional, Union
 
 from ml_orchestrator.artifacts import Artifact
@@ -6,29 +12,94 @@
 
 
 class Model(Artifact):
+    """Artifact type representing a machine learning model.
+
+    This artifact type is used for storing and managing machine learning models
+    in pipeline components.
+
+    Attributes:
+        schema_title: The schema title for this model artifact type
+
+    """
+
     schema_title: str = "system.Model"
 
 
 class Dataset(Artifact):
-    schema_title = "system.Dataset"
+    """Artifact type representing a dataset.
 
-    pass
+    This artifact type is used for storing and managing datasets
+    in pipeline components.
+
+    Attributes:
+        schema_title: The schema title for this dataset artifact type
+
+    """
+
+    schema_title = "system.Dataset"
 
 
 class HTML(Artifact):
+    """Artifact type representing HTML content.
+
+    This artifact type is used for storing and managing HTML content
+    generated by pipeline components, such as visualizations or reports.
+
+    Attributes:
+        schema_title: The schema title for this HTML artifact type
+
+    """
+
     schema_title = "system.HTML"
 
     def __init__(self, name: Optional[str] = None, uri: Optional[str] = None, metadata: Optional[Dict] = None) -> None:
+        """Initialize an HTML artifact with a name, URI, and metadata.
+
+        Ensures the URI ends with .html extension.
+
+        Args:
+            name: The name of the artifact
+            uri: The URI where the artifact is stored
+            metadata: A dictionary of metadata to associate with the artifact
+
+        """
         super().__init__(name=name, uri=uri, metadata=metadata)
         self.uri = self.uri + ".html"
 
 
 class Markdown(Artifact):
+    """Artifact type representing Markdown content.
+
+    This artifact type is used for storing and managing Markdown content
+    generated by pipeline components, such as documentation or reports.
+
+    Attributes:
+        schema_title: The schema title for this Markdown artifact type
+
+    """
+
     schema_title = "system.Markdown"
 
 
 class Metrics(Artifact):
+    """Artifact type representing evaluation metrics.
+
+    This artifact type is used for storing and managing evaluation metrics
+    from pipeline components, such as accuracy, precision, or custom metrics.
+
+    Attributes:
+        schema_title: The schema title for this metrics artifact type
+
+    """
+
     schema_title = "system.Metrics"
 
     def log_metric(self, metric: str, value: MetricTypes) -> None:
+        """Log a metric value with the given name.
+
+        Args:
+            metric: The name of the metric to log
+            value: The value of the metric, which must be one of the supported metric types
+
+        """
         self.metadata[metric] = value

ml_orchestrator/artifacts/notations.py

@@ -0,0 +1,5 @@
+"""Module for artifact type notation definitions.
+
+This module provides type annotations and notations for artifacts
+used in ML pipeline components.
+"""

ml_orchestrator/comp_parser.py

@@ -1,3 +1,10 @@
+"""Module for parsing ML components into Kubeflow Pipeline components.
+
+This module provides a ComponentParser class that extends FunctionParser to create
+and manage components for Kubeflow Pipelines (KFP), handling decorators and
+serialized representations.
+"""
+
 import dataclasses
 from typing import List
 
@@ -8,28 +15,30 @@
 
 @dataclasses.dataclass
 class ComponentParser(FunctionParser):
-    """
-    A parser class that extends FunctionParser to create and manage components
-    for Kubeflow pipelines (KFP), including decorators and serialized representations.
+    """A parser class that extends FunctionParser to create and manage components.
+
+    This class handles components for Kubeflow pipelines (KFP), including decorators
+    and serialized representations.
 
     Attributes:
         add_imports (List[str]): Additional import statements to prepend to output files.
         only_function (bool): Whether to generate only the function body.
+
     """
 
     add_imports: List[str] = dataclasses.field(default_factory=lambda: [])
     only_function: bool = False
 
     @classmethod
     def _create_decorator(cls, env: EnvironmentParams) -> str:
-        """
-        Creates a component decorator string based on environment parameters.
+        """Create a component decorator string based on environment parameters.
 
         Args:
             env (EnvironmentParams): The environment parameters object containing variable definitions.
 
         Returns:
             str: A formatted string representing the component decorator.
+
         """
         dec_vars = env.comp_vars()
         prams = cls.get_func_params(dec_vars, with_typing=False)
@@ -41,28 +50,28 @@ def _create_decorator(cls, env: EnvironmentParams) -> str:
 
     @classmethod
     def create_decorator(cls, component: MetaComponent) -> str:
-        """
-        Creates a decorator string for a given MetaComponent.
+        """Create a decorator string for a given MetaComponent.
 
         Args:
             component (MetaComponent): The component object used to generate a decorator.
 
         Returns:
             str: A decorator string for use in Kubeflow components.
+
         """
         if isinstance(component, MetaComponent):
             return cls._create_decorator(component.env)
         return cls._create_decorator(component.env())
 
     def _create_kfp_str(self, component: _MetaComponent) -> str:  # type: ignore
-        """
-        Generates a serialized Kubeflow Pipeline (KFP) component string.
+        """Generate a serialized Kubeflow Pipeline (KFP) component string.
 
         Args:
             component (_MetaComponent): The component object to serialize.
 
         Returns:
             str: A string representation of the KFP component, including decorators.
+
         """
         function_str = super()._create_kfp_str(component)  # type: ignore
         if self.only_function:
@@ -75,15 +84,15 @@ def _create_kfp_str(self, component: _MetaComponent) -> str:  # type: ignore
         return kfp_component_str + "\n"
 
     def _write_to_file(self, filename: str, file_content: str) -> None:
-        """
-        Writes component string content to a file, including additional imports.
+        """Write component string content to a file, including additional imports.
 
         Args:
             filename (str): The name of the file to write to.
             file_content (str): The content to write to the file.
 
         Returns:
             None
+
         """
         for imp in self.add_imports:
             file_content = f"{imp}\n{file_content}"