Artifact Manager

Overview

The artifact manager subpackage provides a unified interface for logging, retrieving, and managing artifacts across multiple backends. The subpackage defines a base abstract class, ArtifactManager, and includes concrete implementations for the following backends.

FileSystemArtifactManager: Manages artifacts on a filesystem (local or cloud, e.g. using a gs:// bucket).
WandbArtifactManager: Manages artifacts using the Weights & Biases (WandB) platform.

In addition, we also implement convenience context managers: - DirectoryLogger, which is useful for logging an entire folder as an artifact. - ModelLogger, which is a specialized logger for torch models and their states.

Basic Concept

The artifact manager is designed around two core principles: immutability and reproducibility.

Immutability: Each time an artifact is logged, the manager automatically creates a new version using a simple naming convention (for example, “v0”, “v1”, “v2”, etc.). Once an artifact is logged, its version is never modified. This guarantees that past artifacts remain unchanged and can be reliably retrieved, ensuring that experiments or production workflows remain fully reproducible. The artifact manager ensures that every logged artifact is immutable and versioned, enabling you to re-run experiments or roll back to previous states with confidence. You can learn more about artifacts and their use cases from WandB docs.

Unified Interface: By providing a consistent interface regardless of the storage backend (whether a local or cloud-based store or WandB), the artifact manager decouples the logging and retrieval logic from the storage implementation. This makes it easy to switch artifact stores without changing your code. Artifacts are organized into collections (which act as logical groupings) so that different experiments or types of data remain well separated.

Using a Cloud Artifact Store (gs:// Bucket)

The FileSystemArtifactManager can be used with a cloud-based artifact store, such as a Google Cloud Storage bucket. In the example below, we initialize the artifact manager with a gs:// URL, log a file artifact, and download it back to a local directory.

from scaffold.data.artifact_manager.filesystem import FileSystemArtifactManager

# Initialize the artifact manager with a cloud artifact store (gs:// bucket)
manager = FileSystemArtifactManager(url="gs://mybucket/artifacts")

# Log a file artifact.
# This call uploads the file located at "/path/to/local/file.txt" as "my_artifact".
# It also uploads the description as metadata. The exact location is specified in constants.py. Currently, it is "<artifact_base>/meta/readme.txt".
# This (short) description should explain what the artifact is, and for what purpose it was generated.
# log_files() returns an Artifact object with metadata about the logged artifact.
artifact = manager.log_files("my_artifact", "/path/to/local/file.txt", "this is a description that will be logged to a file alongside the artifact")
print(f"Logged artifact: {artifact.name}, version: {artifact.version}, collection: {artifact.collection}")

# Download the artifact to a local directory.
# download_artifact() returns an Artifact object when 'to' is provided.
downloaded_artifact = manager.download_artifact("my_artifact", to="/path/to/download")
print(f"Downloaded artifact: {downloaded_artifact.name}, version: {downloaded_artifact.version}")

# Alternatively, use the context manager to deal with the artifact directly.
# When 'to' is not provided, download_artifact() returns a TmpArtifact context manager
# that has an 'artifact' property with the artifact metadata.
tmp_artifact = manager.download_artifact("my_artifact")
print(f"Temporary artifact: {tmp_artifact.artifact.name}, version: {tmp_artifact.artifact.version}")
with tmp_artifact as tmp_download_dir:
    # Open the file and read its contents, no need to worry about cleanup.
    with open(f"{tmp_download_dir}/file.txt", "r") as f:
        pass

Using the WandB Artifact Manager

The WandbArtifactManager uses the Weights & Biases backend to manage artifacts. Before using this manager, ensure that you have initialized a wandb run (with wandb.init(...)). The example below logs a file artifact and then downloads it.

import wandb
from scaffold.data.artifact_manager.wandb import WandbArtifactManager

# Initialize wandb.
wandb.init(project="my-project")

# Create a WandB artifact manager instance.
manager = WandbArtifactManager(project="my-project", collection="my_collection")

# Log a file artifact.
# log_files() returns an Artifact object with metadata about the logged artifact.
artifact = manager.log_files("example_artifact", "/path/to/local/file.txt", "this is a sample description")
print(f"Logged artifact: {artifact.name}, version: {artifact.version}, collection: {artifact.collection}")

# Download the artifact to a local directory.
# download_artifact() returns an Artifact object when 'to' is provided.
downloaded_artifact = manager.download_artifact("example_artifact", to="/path/to/download")
print(f"Downloaded artifact: {downloaded_artifact.name}, version: {downloaded_artifact.version}")

# Alternatively, use the context manager to deal with the artifact directly.
# When 'to' is not provided, download_artifact() returns a TmpArtifact context manager
# that has an 'artifact' property with the artifact metadata.
tmp_artifact = manager.download_artifact("example_artifact")
print(f"Temporary artifact: {tmp_artifact.artifact.name}, version: {tmp_artifact.artifact.version}")
with tmp_artifact as tmp_download_dir:
    # Open the file and read its contents, no need to worry about cleanup.
    with open(f"{tmp_download_dir}/file.txt", "r") as f:
        pass

Directory Logging Example

You can also log an entire directory as an artifact using the provided DirectoryLogger context manager. In the example below, a temporary directory is created for logging, files are written into it, and then the artifact is automatically logged when the context is exited. Later, the artifact is downloaded and verified.

# Log an entire folder as an artifact.
# log_folder() returns a DirectoryLogger context manager.
logger = manager.log_folder("my_awesome_artifact", "my_collection")
with logger as tmp_dir:
     # Write files to the temporary folder.
     with open(f"{tmp_dir}/example.txt", "w") as f:
         f.write("Example content")
# After the context exits, the artifact is logged and the logger.artifact property
# contains the Artifact object with metadata about the logged artifact.
print(f"Logged artifact: {logger.artifact.name}, version: {logger.artifact.version}")

# Download the logged folder artifact.
# download_artifact() returns an Artifact object when 'to' is provided.
downloaded_artifact = manager.download_artifact("my_awesome_artifact", "my_collection", to="/path/to/download")
print(f"Downloaded artifact: {downloaded_artifact.name}, version: {downloaded_artifact.version}")

# You can also use the context manager:
# When 'to' is not provided, download_artifact() returns a TmpArtifact context manager
# that has an 'artifact' property with the artifact metadata.
tmp_artifact = manager.download_artifact("my_awesome_artifact", "my_collection")
print(f"Temporary artifact: {tmp_artifact.artifact.name}, version: {tmp_artifact.artifact.version}")
with tmp_artifact as tmp_download_dir:
    # Open the file and read its contents, no need to worry about cleanup.
    with open(f"{tmp_download_dir}/example.txt", "r") as f:
        pass

Model Logger

Overview

The ModelLogger class provides a simple interface for logging and retrieving a PyTorch model’s state, including optimizer and scheduler states, as an artifact. The class leverages a temporary local directory to save state dictionaries before logging them to the configured artifact store via an ArtifactManager. This approach enables model state retrieval without needing to reinitialize the model class.

Usage Example: Logging a Model State

The example below demonstrates how to log a model’s state (including optimizer state) as an artifact. In this example, a file system artifact store (using a cloud bucket URL) is used, but any implementation of ArtifactManager (e.g. WandbArtifactManager) can be used instead.

import torch
import torch.nn as nn
import torch.optim as optim
from scaffold.data.artifact_manager.model_logger import ModelLogger
from scaffold.data.artifact_manager.filesystem import FileSystemArtifactManager

# Define a simple model and optimizer.
model = nn.Linear(10, 2)
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Initialize the artifact manager using a cloud artifact store (e.g. a gs:// bucket).
manager = FileSystemArtifactManager(url="gs://mybucket/artifacts")

# Initialize the ModelLogger with the artifact manager.
model_logger = ModelLogger(artifact_manager=manager)

# Log the model state under the artifact id "my_model_state".
afid = model_logger.log_state_to_artifact("my_model_state", model, "this is a sample description that will be logged", optimizers=[optimizer])
print(f"Logged model state with artifact id: {afid}")

Usage Example: Retrieving a Model State

Below is an example of how to retrieve a previously logged model state artifact and load it onto a device.

import torch
from scaffold.data.artifact_manager.model_logger import ModelLogger
from scaffold.data.artifact_manager.filesystem import FileSystemArtifactManager

# Initialize the artifact manager using a cloud artifact store.
manager = FileSystemArtifactManager(url="gs://mybucket/artifacts")

# Initialize the ModelLogger.
model_logger = ModelLogger(artifact_manager=manager)

# Retrieve the model state from the artifact store (e.g., on CPU).
state = model_logger.retrieve_state_from_artifact("my_model_state", device="cpu")
print("Retrieved model state:", state)

Notes

The model state is saved as a file named state.pt within the artifact.
The save_state method unpacks distributed models (if necessary) and saves the unwrapped model along with any optimizer or scheduler state dictionaries.
When retrieving the model state, you can specify the target device (e.g. "cpu" or "cuda") to map the checkpoint appropriately.