How to Implement Custom Auto-Logging for ML Frameworks in CMF

CMF's auto-logging framework in cmflib/contrib/auto_logging_v01.py provides a step decorator and artifact hierarchy that automatically captures execution lineage by wrapping Python functions and dispatching custom artifact types to the appropriate Cmf logging methods.

The Hewlett Packard Enterprise Common Metadata Framework (CMF) offers a lightweight, decorator-based auto-logging system designed to eliminate repetitive instrumentation code in machine learning pipelines. By implementing custom auto-logging for ML frameworks in CMF, you can extend the built-in artifact hierarchy to support proprietary model formats, enabling automatic metadata capture and lineage tracking without modifying core training logic.

Understanding CMF's Auto-Logging Architecture

The auto-logging system centers on three core components implemented in cmflib/contrib/auto_logging_v01.py: the step decorator for execution wrapping, a hierarchical artifact type system, and a central dispatcher that routes artifacts to the appropriate logging methods.

The Step Decorator and Execution Context

The step decorator (lines 251‑405 in auto_logging_v01.py) serves as the primary entry point for auto-logging. When applied to a Python function, it:

• Creates a Cmf instance and initializes a pipeline context and execution
• Introspects function parameters to identify input artifacts (instances of Artifact subclasses)
• Invokes _log_artifacts to persist inputs before function execution
• Executes the wrapped function
• Captures return values and logs them as output artifacts
• Handles the cli_run interface (lines 9‑48) for command-line execution

This decorator effectively transforms any Python function into a tracked pipeline step without requiring manual Cmf method calls.

Artifact Hierarchy and Type System

CMF defines an abstract base class Artifact (lines 80‑133 in auto_logging_v01.py) that standardizes how data assets are represented. Concrete subclasses include:

• Dataset – for training and evaluation data
• MLModel – for standard serialized models
• ExecutionMetrics – for scalar metrics and performance indicators

Each artifact instance carries a uri (file path), params (metadata dictionary), and type information. When you implement custom auto-logging for a new ML framework, you create a new subclass (e.g., MyFancyModel) that inherits from Artifact, instantly making it compatible with the step decorator's parameter inspection.

The Dispatcher Pattern in _log_artifacts

The _log_artifacts function (lines 442‑472 in auto_logging_v01.py) acts as a central dispatcher that maps artifact types to specific Cmf logging methods. The current implementation handles built-in types:

# Simplified logic from auto_logging_v01.py lines 442-472

if isinstance(artifacts, Dataset):
    cmf.log_dataset(path=artifacts.uri, event=event, **artifacts.params)
elif isinstance(artifacts, MLModel):
    cmf.log_model(path=artifacts.uri, event=event, **artifacts.params)
elif isinstance(artifacts, ExecutionMetrics):
    cmf.log_execution_metrics(path=artifacts.uri, **artifacts.params)

To support a custom artifact type, you extend this dispatcher with a new elif branch that invokes your custom logging method.

Step-by-Step Implementation Guide

Extending CMF's auto-logging to support a custom ML framework requires five specific modifications across cmflib/contrib/auto_logging_v01.py and cmflib/cmf.py.

Step 1: Define a Custom Artifact Subclass

Create a new class inheriting from Artifact to represent your framework's specific model format. Add this to cmflib/contrib/auto_logging_v01.py near the existing artifact definitions (lines 80‑133):

class MyFancyModel(Artifact):
    """Artifact representing a MyFancyNet model checkpoint."""
    def __init__(self, uri: str, params: dict = None):
        super().__init__(uri, params)
        self.type_name = "MyFancyModel"

This class standardizes how your custom model is packaged for the auto-logging system.

Step 2: Implement the Cmf Logging Method

Add a dedicated logging method to the Cmf class in cmflib/cmf.py, following the pattern established by log_model (around line 510). This method handles the low-level MLMD artifact creation and DVC integration:

def log_my_fancy_model(self, path: str, event: str, **custom_props):
    """Log a MyFancyNet model as a CMF artifact.
    
    Args:
        path: Filesystem path to the model checkpoint
        event: 'input' or 'output' to set artifact lineage direction
        **custom_props: Framework-specific metadata (version, accuracy, etc.)
    """
    import time
    from cmflib import mlpb
    
    # Commit file with DVC and obtain content hash

    commit_output(path, self.execution.id)
    c_hash = dvc_get_hash(path)
    
    # Build artifact URI with hash for immutability

    uri = f"{path}:{c_hash}"
    
    # Create MLMD artifact event

    artifact = create_new_artifact_event_and_attribution(
        store=self.store,
        execution_id=self.execution.id,
        context_id=self.child_context.id,
        uri=c_hash,
        name=uri,
        type_name="MyFancyModel",
        event_type=mlpb.Event.Type.OUTPUT if event == "output" else mlpb.Event.Type.INPUT,
        properties=custom_props,
        artifact_type_properties={k: mlpb.STRING for k in custom_props},
        milliseconds_since_epoch=int(time.time() * 1000),
    )
    return artifact

This method ensures your custom model type integrates with CMF's DVC-backed storage and MLMD lineage tracking.

Step 3: Extend the Artifact Dispatcher

Modify _log_artifacts in cmflib/contrib/auto_logging_v01.py (lines 442‑472) to recognize MyFancyModel and route it to your new logging method:

elif isinstance(artifacts, MyFancyModel):
    cmf.log_my_fancy_model(path=artifacts.uri, event=event, **artifacts.params)

Insert this branch alongside the existing Dataset, MLModel, and ExecutionMetrics checks to ensure seamless integration.

Step 4: Use the Custom Artifact in Pipeline Steps

With the infrastructure in place, use your custom artifact type in pipeline functions decorated with @step(). Create a training script at examples/auto_logging_v01/pipeline/train_fancy.py:

from cmflib.contrib.auto_logging_v01 import (
    step, Context, Parameters, Dataset, MyFancyModel, cli_run, prepare_workspace
)

@step()
def train_fancy(ctx: Context, params: Parameters, train_set: Dataset) -> dict[str, MyFancyModel]:
    """Train a MyFancyNet model with automatic CMF logging."""
    # Initialize workspace

    workspace = prepare_workspace(ctx)
    model_path = workspace / "fancy_model.pkl"
    
    # Execute framework-specific training (outside CMF)

    from my_fancy_lib import train_fancy_net
    train_fancy_net(input_path=train_set.uri, output_path=model_path)
    
    # Return custom artifact; CMF auto-logs via the dispatcher

    return {"model": MyFancyModel(model_path, {"accuracy": 0.93, "framework_version": "2.1.0"})}

if __name__ == "__main__":
    cli_run(train_fancy)

When executed via the command line:

python examples/auto_logging_v01/pipeline/train_fancy.py \
    --ctx workspace=workspace \
    train_set=workspace/iris.pkl

CMF automatically creates a pipeline execution, logs the input Dataset, and persists the MyFancyModel output with full lineage metadata.

Step 5: Expose the New Artifact in the Public API

Add MyFancyModel to the __all__ list at the top of cmflib/contrib/auto_logging_v01.py to make it available for direct import:

__all__ = [
    "step",
    "Context", 
    "Parameters",
    "Dataset",
    "MLModel",
    "ExecutionMetrics",
    "MyFancyModel",  # Custom artifact exposed

    "cli_run",
    "prepare_workspace"
]

This completes the integration, allowing users to import your custom artifact type alongside built-in CMF components.

Key Source Files and Their Roles

Understanding the repository structure is essential for implementing custom auto-logging:

• cmflib/contrib/auto_logging_v01.py – Implements the public auto-logging API including the step decorator (lines 251‑405), artifact class hierarchy (lines 80‑133), _log_artifacts dispatcher (lines 442‑472), and cli_run interface (lines 9‑48).

• cmflib/cmf.py – Contains the core Cmf client class that creates contexts, executions, and provides low-level log_* helpers (such as log_model around line 510) used by the dispatcher.

• examples/auto_logging_v01/pipeline/train.py – Reference implementation demonstrating standard MLModel auto-logging patterns.

• examples/auto_logging_v01/pipeline/preprocess.py – Demonstrates Dataset artifact handling and workspace preparation utilities.

Summary

Implementing custom auto-logging for ML frameworks in CMF requires extending three architectural layers:

• Artifact Types – Create subclasses of Artifact in auto_logging_v01.py to represent custom model formats
• Logging Methods – Implement framework-specific log_* methods in cmf.py to handle DVC hashing and MLMD artifact creation
• Dispatch Logic – Extend _log_artifacts to route custom types to their respective logging methods

Once registered, the step decorator automatically captures lineage for your custom artifacts without requiring manual instrumentation, enabling consistent metadata capture and graph-level lineage tracking across heterogeneous ML toolchains.

Frequently Asked Questions

What is the CMF auto-logging framework?

The CMF auto-logging framework is a decorator-based system implemented in cmflib/contrib/auto_logging_v01.py that automatically instruments Python functions to capture machine learning metadata. It wraps user functions with the @step() decorator to create MLMD executions, log input artifacts before function invocation, and capture output artifacts upon completion, eliminating the need for repetitive Cmf method calls throughout pipeline code.

How does the step decorator work in CMF?

The step decorator (lines 251‑405 in auto_logging_v01.py) functions as a higher-order function that intercepts pipeline function calls. It creates a Cmf instance and initializes a pipeline context and execution, introspects function parameters to identify input artifacts, invokes _log_artifacts to persist inputs before execution, runs the wrapped function, and processes return values through the same dispatcher to log outputs with full lineage metadata.

Can I auto-log custom model formats beyond standard MLModel?

Yes, CMF supports custom model formats through a three-part extension mechanism. First, create a new subclass of Artifact (e.g., MyFancyModel) in auto_logging_v01.py. Second, implement a corresponding log_* method in cmflib/cmf.py that handles DVC hashing and MLMD artifact creation. Third, extend the _log_artifacts dispatcher (lines 442‑472) to route your custom type to the new logging method, enabling automatic capture without manual instrumentation.

Where does CMF store the metadata for auto-logged artifacts?

CMF stores auto-logged metadata in an MLMD (ML Metadata) database, typically configured as a SQLite file specified during Cmf initialization. When DVC integration is enabled, CMF commits output files to DVC storage and stores the resulting content hash in MLMD as the artifact URI, ensuring immutable versioning. If the graph=True parameter is set during initialization, CMF additionally synchronizes metadata to a Neo4j graph database for visual lineage exploration.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}