Dagster

Integrate W&B with Dagster to track ML experiments and manage data pipelines with automatic logging and monitoring.

Use Dagster and W&B together to orchestrate your MLOps pipelines and maintain ML assets, so you can track experiments, manage data, and run training jobs from within your existing Dagster workflows. The integration with W&B lets you do the following within Dagster:

• Create and use a W&B Artifact.
• Use and create Registered Models in W&B Registry.
• Run training jobs on dedicated compute using W&B Launch.
• Use the wandb client in ops and assets.

The W&B Dagster integration provides a W&B-specific Dagster resource and IO Manager:

• wandb_resource: a Dagster resource used to authenticate and communicate with the W&B API.
• wandb_artifacts_io_manager: a Dagster IO Manager used to consume W&B Artifacts.

This guide is for ML practitioners and platform engineers who already use Dagster and want to add W&B tracking and artifact management. It walks through the prerequisites for using W&B in Dagster, shows how to create and use W&B Artifacts in ops and assets, explains how to use W&B Launch, and describes recommended best practices.

Before you get started

Before you configure the integration, gather the W&B credentials and identifiers it needs to authenticate and route your runs.

You need the following resources to use Dagster within W&B:

• W&B API Key.
• W&B entity (user or team): An entity is a username or team name where you send W&B Runs and Artifacts. Create your account or team entity in the W&B App UI before you log runs. If you don’t specify an entity, the run goes to your default entity, which is usually your username. Change your default entity in your settings under Project Defaults.
• W&B project: The name of the project where W&B Runs are stored.

Find your W&B entity by checking the profile page for that user or team in the W&B App. You can use a pre-existing W&B project or create a new one. Create new projects on the W&B App homepage or on the user or team profile page. If a project doesn’t exist, W&B creates it automatically when you first use it.

Set up your API key

The integration authenticates with the W&B API using an API key, which you must make available to Dagster as an environment variable.

1. Log in to W&B. If you’re using W&B Server, ask your admin for the instance host name.
2. Create an API key at User Settings. For a production environment, use a service account to own that key.
3. Set an environment variable for that API key: export WANDB_API_KEY=[YOUR_KEY].

After completing these steps, Dagster can read your API key from the environment when it loads the wandb_resource.

The following examples demonstrate where to specify your API key in your Dagster code. Specify your entity and project name within the wandb_config nested dictionary. You can pass different wandb_config values to different ops or assets if you want to use a different W&B Project. For more information about possible keys you can pass, see the following Configuration section.

Config for @job

Example configuration for @job:

# add this to your config.yaml
# alternatively you can set the config in Dagit's Launchpad or JobDefinition.execute_in_process
# Reference: https://docs.dagster.io/concepts/configuration/config-schema#specifying-runtime-configuration
resources:
 wandb_config:
   config:
     entity: my_entity # replace this with your W&B entity
     project: my_project # replace this with your W&B project


@job(
   resource_defs={
       "wandb_config": make_values_resource(
           entity=str,
           project=str,
       ),
       "wandb_resource": wandb_resource.configured(
           {"api_key": {"env": "WANDB_API_KEY"}}
       ),
       "io_manager": wandb_artifacts_io_manager,
   }
)
def simple_job_example():
   my_op()

Config for @repository using assets

Example configuration for @repository using assets:

from dagster_wandb import wandb_artifacts_io_manager, wandb_resource
from dagster import (
   load_assets_from_package_module,
   make_values_resource,
   repository,
   with_resources,
)

from . import assets

@repository
def my_repository():
   return [
       *with_resources(
           load_assets_from_package_module(assets),
           resource_defs={
               "wandb_config": make_values_resource(
                   entity=str,
                   project=str,
               ),
               "wandb_resource": wandb_resource.configured(
                   {"api_key": {"env": "WANDB_API_KEY"}}
               ),
               "wandb_artifacts_manager": wandb_artifacts_io_manager.configured(
                   {"cache_duration_in_minutes": 60} # only cache files for one hour
               ),
           },
           resource_config_by_key={
               "wandb_config": {
                   "config": {
                       "entity": "my_entity", # replace this with your W&B entity
                       "project": "my_project", # replace this with your W&B project
                   }
               }
           },
       ),
   ]

This example configures the IO Manager cache duration, unlike the example for @job.

Configuration

The integration provides a W&B-specific Dagster resource and IO Manager with the following configuration options.

• wandb_resource: Dagster resource used to communicate with the W&B API. It authenticates automatically using the provided API key. Properties:
  • api_key: (str, required): a W&B API key necessary to communicate with the W&B API.
  • host: (str, optional): the API host server you want to use. Only required if you’re using W&B Server. Defaults to the Public Cloud host, https://api.wandb.ai.
• wandb_artifacts_io_manager: Dagster IO Manager to consume W&B Artifacts. Properties:
  • base_dir: (int, optional) Base directory used for local storage and caching. W&B Artifacts and W&B Run logs are written to and read from that directory. By default, it uses the DAGSTER_HOME directory.
  • cache_duration_in_minutes: (int, optional) Defines how long to keep W&B Artifacts and W&B Run logs in local storage. The cache removes only files and directories that haven’t been opened for that amount of time. Cache purging happens at the end of an IO Manager execution. Set it to 0 to turn off caching completely. Caching improves speed when an Artifact is reused between jobs running on the same machine. Defaults to 30 days.
  • run_id: (str, optional): A unique ID for this run, used for resuming. It must be unique in the project, and if you delete a run you can’t reuse the ID. Use the name field for a short descriptive name, or config for saving hyperparameters to compare across runs. The ID can’t contain the following special characters: /\#?%:.. Set the Run ID when you’re doing experiment tracking inside Dagster to allow the IO Manager to resume the run. By default, it’s set to the Dagster Run ID, for example, 7e4df022-1bf2-44b5-a383-bb852df4077e.
  • run_name: (str, optional) A short display name for this run to help you identify it in the UI. By default, it is a string with the following format: dagster-run-[8 first characters of the Dagster Run ID]. For example, dagster-run-7e4df022.
  • run_tags: (list[str], optional): A list of strings that populates the list of tags on this run in the UI. Tags are useful for organizing runs together or applying temporary labels like baseline or production. You can add and remove tags in the UI, or filter down to runs with a specific tag. Any W&B Run used by the integration has the dagster_wandb tag.

Use W&B Artifacts

This section explains how the integration uses a Dagster IO Manager to bridge W&B Artifacts and Dagster ops and assets.

The integration with W&B Artifact relies on a Dagster IO Manager.

IO Managers are user-provided objects that store the output of an asset or op and load it as input to downstream assets or ops. For example, an IO Manager might store and load objects from files on a filesystem.

The integration provides an IO Manager for W&B Artifacts. This lets any Dagster @op or @asset create and consume W&B Artifacts directly. The following example shows an @asset that produces a W&B Artifact of type dataset containing a Python list.

@asset(
    name="my_artifact",
    metadata={
        "wandb_artifact_arguments": {
            "type": "dataset",
        }
    },
    io_manager_key="wandb_artifacts_manager",
)
def create_dataset():
    return [1, 2, 3] # this will be stored in an Artifact

You can annotate your @op, @asset, and @multi_asset with a metadata configuration to write Artifacts. Similarly, you can also consume W&B Artifacts even if they were created outside Dagster.

Write W&B Artifacts

The following sections describe how to produce W&B Artifacts from Dagster ops and assets, including the supported return types and how to configure the resulting Artifact.

Before continuing, make sure you understand how to use W&B Artifacts. See the Guide on Artifacts.

Return an object from a Python function to write a W&B Artifact. W&B supports the following objects:

• Python objects (int, dict, list, and so on)
• W&B objects (Table, Image, Graph, and so on)
• W&B Artifact objects

The following examples demonstrate how to write W&B Artifacts with Dagster assets (@asset):

Python objects

Anything that can be serialized with the pickle module is pickled and added to an Artifact created by the integration. The content is unpickled when you read that Artifact inside Dagster (see Read artifacts for more details).

@asset(
    name="my_artifact",
    metadata={
        "wandb_artifact_arguments": {
            "type": "dataset",
        }
    },
    io_manager_key="wandb_artifacts_manager",
)
def create_dataset():
    return [1, 2, 3]

W&B supports multiple Pickle-based serialization modules (pickle, dill, cloudpickle, joblib). You can also use more advanced serialization like ONNX or PMML. Refer to the Serialization section for more information.

W&B Object

Any W&B object, such as a Table or Image, is added to an Artifact created by the integration. This example adds a Table to an Artifact:

import wandb

@asset(
    name="my_artifact",
    metadata={
        "wandb_artifact_arguments": {
            "type": "dataset",
        }
    },
    io_manager_key="wandb_artifacts_manager",
)
def create_dataset_in_table():
    return wandb.Table(columns=["a", "b", "c"], data=[[1, 2, 3]])

W&B Artifact

For complex use cases, you might need to build your own Artifact object. The integration still provides additional features like augmenting the metadata on both sides of the integration.

import wandb

MY_ASSET = "my_asset"

@asset(
    name=MY_ASSET,
    io_manager_key="wandb_artifacts_manager",
)
def create_artifact():
   artifact = wandb.Artifact(MY_ASSET, "dataset")
   table = wandb.Table(columns=["a", "b", "c"], data=[[1, 2, 3]])
   artifact.add(table, "my_table")
   return artifact

Configuration

You can set a configuration dictionary called wandb_artifact_configuration on an @op, @asset, and @multi_asset. Pass this dictionary in the decorator arguments as metadata. This configuration is required to control the IO Manager reads and writes of W&B Artifacts.

For @op, it’s located in the output metadata through the Out metadata argument. For @asset, it’s located in the metadata argument on the asset. For @multi_asset, it’s located in each output metadata through the AssetOut metadata arguments.

The following code examples demonstrate how to configure a dictionary on an @op, @asset, and @multi_asset computations:

Example for @op

Example for @op:

@op(
   out=Out(
       metadata={
           "wandb_artifact_configuration": {
               "name": "my_artifact",
               "type": "dataset",
           }
       }
   )
)
def create_dataset():
   return [1, 2, 3]

Example for @asset

Example for @asset:

@asset(
   name="my_artifact",
   metadata={
       "wandb_artifact_configuration": {
           "type": "dataset",
       }
   },
   io_manager_key="wandb_artifacts_manager",
)
def create_dataset():
   return [1, 2, 3]

You don’t need to pass a name through the configuration because the @asset already has a name. The integration sets the Artifact name as the asset name.

Example for @multi_asset

Example for @multi_asset:

@multi_asset(
   name="create_datasets",
   outs={
       "first_table": AssetOut(
           metadata={
               "wandb_artifact_configuration": {
                   "type": "training_dataset",
               }
           },
           io_manager_key="wandb_artifacts_manager",
       ),
       "second_table": AssetOut(
           metadata={
               "wandb_artifact_configuration": {
                   "type": "validation_dataset",
               }
           },
           io_manager_key="wandb_artifacts_manager",
       ),
   },
   group_name="my_multi_asset_group",
)
def create_datasets():
   first_table = wandb.Table(columns=["a", "b", "c"], data=[[1, 2, 3]])
   second_table = wandb.Table(columns=["d", "e"], data=[[4, 5]])

   return first_table, second_table

The following properties are supported:

• name: (str) A human-readable name for this artifact that you can use to identify it in the UI or reference it in use_artifact calls. Names can contain letters, numbers, underscores, hyphens, and dots. The name must be unique across a project. Required for @op.
• type: (str) The type of the artifact, used to organize and differentiate artifacts. Common types include dataset or model, but you can use any string containing letters, numbers, underscores, hyphens, and dots. Required when the output isn’t already an Artifact.
• description: (str) Free text that describes the artifact. The description is rendered as markdown in the UI, so it’s a good place for tables, links, and so on.
• aliases: (list[str]) An array containing one or more aliases you want to apply on the Artifact. The integration also adds the “latest” tag to that list whether it’s set or not. Use aliases to manage versioning of models and datasets.
• add_dirs: (list[dict[str, Any]]): An array containing configuration for each local directory to include in the Artifact.
• add_files: (list[dict[str, Any]]): An array containing configuration for each local file to include in the Artifact.
• add_references: (list[dict[str, Any]]): An array containing configuration for each external reference to include in the Artifact.
• serialization_module: (dict) Configuration of the serialization module to use. Refer to the Serialization section for more information.
  • name: (str) Name of the serialization module. Accepted values: pickle, dill, cloudpickle, joblib. The module must be available locally.
  • parameters: (dict[str, Any]) Optional arguments passed to the serialization function. It accepts the same parameters as the dump method for that module. For example, {"compress": 3, "protocol": 4}.

Advanced example:

@asset(
   name="my_advanced_artifact",
   metadata={
       "wandb_artifact_configuration": {
           "type": "dataset",
           "description": "My *Markdown* description",
           "aliases": ["my_first_alias", "my_second_alias"],
           "add_dirs": [
               {
                   "name": "My directory",
                   "local_path": "path/to/directory",
               }
           ],
           "add_files": [
               {
                   "name": "validation_dataset",
                   "local_path": "path/to/data.json",
               },
               {
                   "is_tmp": True,
                   "local_path": "path/to/temp",
               },
           ],
           "add_references": [
               {
                   "uri": "https://picsum.photos/200/300",
                   "name": "External HTTP reference to an image",
               },
               {
                   "uri": "s3://my-bucket/datasets/mnist",
                   "name": "External S3 reference",
               },
           ],
       }
   },
   io_manager_key="wandb_artifacts_manager",
)
def create_advanced_artifact():
   return [1, 2, 3]

The asset is materialized with metadata on both sides of the integration:

• W&B side: the source integration name and version, the Python version used, the pickle protocol version, and more.
• Dagster side:
  • Dagster Run ID
  • W&B Run: ID, name, path, URL
  • W&B Artifact: ID, name, type, version, size, URL
  • W&B Entity
  • W&B Project

The following image demonstrates the metadata from W&B that the integration adds to the Dagster asset. The integration propagates this information to Dagster.

The following image demonstrates how the provided configuration is enriched with metadata on the W&B Artifact. This information helps with reproducibility and maintenance. It isn’t available without the integration.

If you use a static type checker like mypy, import the configuration type definition object using the following:

from dagster_wandb import WandbArtifactConfiguration

Use partitions

The integration directly supports Dagster partitions.

The following is an example with a partitioned using DailyPartitionsDefinition:

@asset(
    partitions_def=DailyPartitionsDefinition(start_date="2023-01-01", end_date="2023-02-01"),
    name="my_daily_partitioned_asset",
    compute_kind="wandb",
    metadata={
        "wandb_artifact_configuration": {
            "type": "dataset",
        }
    },
)
def create_my_daily_partitioned_asset(context):
    partition_key = context.asset_partition_key_for_output()
    context.log.info(f"Creating partitioned asset for {partition_key}")
    return random.randint(0, 100)

This code produces one W&B Artifact for each partition. View artifacts in the Artifact panel (UI) under the asset name, which has the partition key appended. For example, my_daily_partitioned_asset.2023-01-01, my_daily_partitioned_asset.2023-01-02, or my_daily_partitioned_asset.2023-01-03. Assets that are partitioned across multiple dimensions show each dimension in dot-delimited format. For example, my_asset.car.blue.

The integration doesn’t allow for the materialization of multiple partitions within one run. You need to perform multiple runs to materialize your assets. You can execute this in Dagit when you’re materializing your assets.

Advanced usage

For advanced usage, see the following examples:

• Partitioned job
• Simple partitioned asset
• Multi-partitioned asset
• Advanced partitioned usage

Read W&B Artifacts

Now that you can write Artifacts from Dagster, the following sections describe how to consume them as inputs to downstream ops and assets.

Reading W&B Artifacts is similar to writing them. You can set a configuration dictionary called wandb_artifact_configuration on an @op or @asset. The only difference is that you set the configuration on the input instead of the output.

For @op, it’s located in the input metadata through the In metadata argument. You need to explicitly pass the name of the Artifact.

For @asset, it’s located in the input metadata through the Asset In metadata argument. Don’t pass an Artifact name because the name of the parent asset should match it.

If you want a dependency on an Artifact created outside the integration, use SourceAsset. It always reads the latest version of that asset.

The following examples demonstrate how to read an Artifact from various ops.

From an @op

Reading an artifact from an @op:

@op(
   ins={
       "artifact": In(
           metadata={
               "wandb_artifact_configuration": {
                   "name": "my_artifact",
               }
           }
       )
   },
   io_manager_key="wandb_artifacts_manager"
)
def read_artifact(context, artifact):
   context.log.info(artifact)

Created by another @asset

Reading an artifact created by another @asset:

@asset(
   name="my_asset",
   ins={
       "artifact": AssetIn(
           # if you don't want to rename the input argument you can remove 'key'
           key="parent_dagster_asset_name",
           input_manager_key="wandb_artifacts_manager",
       )
   },
)
def read_artifact(context, artifact):
   context.log.info(artifact)

Artifact created outside Dagster

Reading an Artifact created outside Dagster:

my_artifact = SourceAsset(
   key=AssetKey("my_artifact"),  # the name of the W&B Artifact
   description="Artifact created outside Dagster",
   io_manager_key="wandb_artifacts_manager",
)


@asset
def read_artifact(context, my_artifact):
   context.log.info(my_artifact)

Configuration

The following configuration indicates what the IO Manager should collect and provide as inputs to the decorated functions. The following read patterns are supported:

• To get a named object contained within an Artifact, use get:

@asset(
   ins={
       "table": AssetIn(
           key="my_artifact_with_table",
           metadata={
               "wandb_artifact_configuration": {
                   "get": "my_table",
               }
           },
           input_manager_key="wandb_artifacts_manager",
       )
   }
)
def get_table(context, table):
   context.log.info(table.get_column("a"))

• To get the local path of a downloaded file contained within an Artifact, use get_path:

@asset(
   ins={
       "path": AssetIn(
           key="my_artifact_with_file",
           metadata={
               "wandb_artifact_configuration": {
                   "get_path": "name_of_file",
               }
           },
           input_manager_key="wandb_artifacts_manager",
       )
   }
)
def get_path(context, path):
   context.log.info(path)

• To get the entire Artifact object (with the content downloaded locally):

@asset(
   ins={
       "artifact": AssetIn(
           key="my_artifact",
           input_manager_key="wandb_artifacts_manager",
       )
   },
)
def get_artifact(context, artifact):
   context.log.info(artifact.name)

The following properties are supported:

• get: (str) Gets the W&B object located at the artifact relative name.
• get_path: (str) Gets the path to the file located at the artifact relative name.

Serialization configuration

By default, the integration uses the standard pickle module, but some objects aren’t compatible with it. For example, functions with yield raise an error if you try to pickle them.

W&B supports more Pickle-based serialization modules (dill, cloudpickle, joblib). You can also use more advanced serialization like ONNX or PMML by returning a serialized string or creating an Artifact directly. The right choice depends on your use case. Refer to the available literature on this subject.

Pickle-based serialization modules

Pickling is known to be insecure. If security is a concern, use only W&B objects. Sign your data and store the hash keys in your own systems. For more complex use cases, contact W&B Support.

You can configure the serialization used through the serialization_module dictionary in the wandb_artifact_configuration. Make sure the module is available on the machine running Dagster.

The integration automatically detects which serialization module to use when you read that Artifact.

The supported modules are pickle, dill, cloudpickle, and joblib.

Here’s a simplified example that creates a “model” serialized with joblib and then uses it for inference.

@asset(
    name="my_joblib_serialized_model",
    compute_kind="Python",
    metadata={
        "wandb_artifact_configuration": {
            "type": "model",
            "serialization_module": {
                "name": "joblib"
            },
        }
    },
    io_manager_key="wandb_artifacts_manager",
)
def create_model_serialized_with_joblib():
    # This is not a real ML model but this would not be possible with the pickle module
    return lambda x, y: x + y

@asset(
    name="inference_result_from_joblib_serialized_model",
    compute_kind="Python",
    ins={
        "my_joblib_serialized_model": AssetIn(
            input_manager_key="wandb_artifacts_manager",
        )
    },
    metadata={
        "wandb_artifact_configuration": {
            "type": "results",
        }
    },
    io_manager_key="wandb_artifacts_manager",
)
def use_model_serialized_with_joblib(
    context: OpExecutionContext, my_joblib_serialized_model
):
    inference_result = my_joblib_serialized_model(1, 2)
    context.log.info(inference_result)  # Prints: 3
    return inference_result

Advanced serialization formats (ONNX, PMML)

Interchange file formats like ONNX and PMML are common. The integration supports those formats, but it requires a bit more work than Pickle-based serialization.

You can use these formats with one of the following methods:

• Convert your model to the selected format, then return the string representation of that format as if it were a normal Python object. The integration pickles that string. You can then rebuild your model using that string.
• Create a new local file with your serialized model, then build a custom Artifact with that file using the add_file configuration.

Here’s an example of a Scikit-learn model serialized using ONNX.

import numpy
import onnxruntime as rt
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split

from dagster import AssetIn, AssetOut, asset, multi_asset

@multi_asset(
    compute_kind="Python",
    outs={
        "my_onnx_model": AssetOut(
            metadata={
                "wandb_artifact_configuration": {
                    "type": "model",
                }
            },
            io_manager_key="wandb_artifacts_manager",
        ),
        "my_test_set": AssetOut(
            metadata={
                "wandb_artifact_configuration": {
                    "type": "test_set",
                }
            },
            io_manager_key="wandb_artifacts_manager",
        ),
    },
    group_name="onnx_example",
)
def create_onnx_model():
    # Inspired from https://onnx.ai/sklearn-onnx/

    # Train a model.
    iris = load_iris()
    X, y = iris.data, iris.target
    X_train, X_test, y_train, y_test = train_test_split(X, y)
    clr = RandomForestClassifier()
    clr.fit(X_train, y_train)

    # Convert into ONNX format
    initial_type = [("float_input", FloatTensorType([None, 4]))]
    onx = convert_sklearn(clr, initial_types=initial_type)

    # Write artifacts (model + test_set)
    return onx.SerializeToString(), {"X_test": X_test, "y_test": y_test}

@asset(
    name="experiment_results",
    compute_kind="Python",
    ins={
        "my_onnx_model": AssetIn(
            input_manager_key="wandb_artifacts_manager",
        ),
        "my_test_set": AssetIn(
            input_manager_key="wandb_artifacts_manager",
        ),
    },
    group_name="onnx_example",
)
def use_onnx_model(context, my_onnx_model, my_test_set):
    # Inspired from https://onnx.ai/sklearn-onnx/

    # Compute the prediction with ONNX Runtime
    sess = rt.InferenceSession(my_onnx_model)
    input_name = sess.get_inputs()[0].name
    label_name = sess.get_outputs()[0].name
    pred_onx = sess.run(
        [label_name], {input_name: my_test_set["X_test"].astype(numpy.float32)}
    )[0]
    context.log.info(pred_onx)
    return pred_onx

Use partitions

The integration directly supports Dagster partitions.

You can selectively read one, multiple, or all partitions of an asset.

All partitions are provided in a dictionary, with the key and value representing the partition key and the Artifact content, respectively.

Read all partitions

This reads all partitions of the upstream @asset, which are given as a dictionary. In this dictionary, the key and value correlate to the partition key and the Artifact content, respectively.

@asset(
    compute_kind="wandb",
    ins={"my_daily_partitioned_asset": AssetIn()},
    output_required=False,
)
def read_all_partitions(context, my_daily_partitioned_asset):
    for partition, content in my_daily_partitioned_asset.items():
        context.log.info(f"partition={partition}, content={content}")

Read specific partitions

The AssetIn’s partition_mapping configuration lets you choose specific partitions. In this case, the example uses the TimeWindowPartitionMapping.

@asset(
    partitions_def=DailyPartitionsDefinition(start_date="2023-01-01", end_date="2023-02-01"),
    compute_kind="wandb",
    ins={
        "my_daily_partitioned_asset": AssetIn(
            partition_mapping=TimeWindowPartitionMapping(start_offset=-1)
        )
    },
    output_required=False,
)
def read_specific_partitions(context, my_daily_partitioned_asset):
    for partition, content in my_daily_partitioned_asset.items():
        context.log.info(f"partition={partition}, content={content}")

The configuration object, metadata, configures how W&B interacts with different artifact partitions in your project.

The object metadata contains a key named wandb_artifact_configuration which further contains a nested object partitions.

The partitions object maps the name of each partition to its configuration. The configuration for each partition can specify how to retrieve data from it. These configurations can contain different keys, namely get, version, and alias, depending on the requirements of each partition.

Configuration keys

The following configuration keys are supported:

• get: The get key specifies the name of the W&B Object (Table, Image, and so on) to fetch the data from.
• version: Use the version key when you want to fetch a specific version of the Artifact.
• alias: The alias key lets you get the Artifact by its alias.

Wildcard configuration

The wildcard "*" stands for all non-configured partitions. This provides a default configuration for partitions that aren’t explicitly mentioned in the partitions object.

For example:

"*": {
    "get": "default_table_name",
},

This configuration means that for all partitions not explicitly configured, data is fetched from the table named default_table_name.

Specific partition configuration

You can override the wildcard configuration for specific partitions by providing their specific configurations using their keys.

For example:

"yellow": {
    "get": "custom_table_name",
},

This configuration means that for the partition named yellow, data is fetched from the table named custom_table_name, overriding the wildcard configuration.

Versioning and aliasing

For versioning and aliasing purposes, you can provide specific version and alias keys in your configuration.

For versions:

"orange": {
    "version": "v0",
},

This configuration fetches data from version v0 of the orange Artifact partition.

For aliases:

"blue": {
    "alias": "special_alias",
},

This configuration fetches data from the table default_table_name of the Artifact partition with the alias special_alias (referred to as blue in the configuration).

Advanced usage

To view advanced usage of the integration, refer to the following full code examples:

• Advanced usage example for assets
• Partitioned job example
• Linking a model to the Model Registry

Use W&B Launch

The following sections describe how to combine the Dagster integration with W&B Launch to run training jobs on dedicated compute, either locally or remotely.

Beta product in active development. Interested in Launch? Reach out to your account team to talk about joining the customer pilot program for W&B Launch. Pilot customers need to use AWS EKS or SageMaker to qualify for the beta program. Additional platforms are planned.

Before continuing, make sure you understand how to use W&B Launch. See the Guide on Launch.

The Dagster integration helps with:

• Running one or multiple Launch agents in your Dagster instance.
• Executing local Launch jobs within your Dagster instance.
• Remote Launch jobs on-prem or in a cloud.

Launch agents

The integration provides an importable @op called run_launch_agent. It starts a Launch Agent and runs it as a long-running process until stopped manually.

Agents are processes that poll launch queues and execute the jobs (or dispatch them to external services to be executed) in order.

Refer to the Launch page.

You can also view descriptions for all properties in Launchpad.

Example:

# add this to your config.yaml
# alternatively you can set the config in Dagit's Launchpad or JobDefinition.execute_in_process
# Reference: https://docs.dagster.io/concepts/configuration/config-schema#specifying-runtime-configuration
resources:
 wandb_config:
   config:
     entity: my_entity # replace this with your W&B entity
     project: my_project # replace this with your W&B project
ops:
 run_launch_agent:
   config:
     max_jobs: -1
     queues: 
       - my_dagster_queue

from dagster_wandb.launch.ops import run_launch_agent
from dagster_wandb.resources import wandb_resource

from dagster import job, make_values_resource

@job(
   resource_defs={
       "wandb_config": make_values_resource(
           entity=str,
           project=str,
       ),
       "wandb_resource": wandb_resource.configured(
           {"api_key": {"env": "WANDB_API_KEY"}}
       ),
   },
)
def run_launch_agent_example():
   run_launch_agent()

Launch jobs

The integration provides an importable @op called run_launch_job. It executes your Launch job.

A Launch job is assigned to a queue to be executed. You can create a queue or use the default one. Make sure you have an active agent listening to that queue. You can run an agent inside your Dagster instance, or consider using a deployable agent in Kubernetes.

Refer to the Launch page.

You can also view descriptions for all properties in Launchpad.

Example:

# add this to your config.yaml
# alternatively you can set the config in Dagit's Launchpad or JobDefinition.execute_in_process
# Reference: https://docs.dagster.io/concepts/configuration/config-schema#specifying-runtime-configuration
resources:
 wandb_config:
   config:
     entity: my_entity # replace this with your W&B entity
     project: my_project # replace this with your W&B project
ops:
 my_launched_job:
   config:
     entry_point:
       - python
       - train.py
     queue: my_dagster_queue
     uri: https://github.com/wandb/example-dagster-integration-with-launch


from dagster_wandb.launch.ops import run_launch_job
from dagster_wandb.resources import wandb_resource

from dagster import job, make_values_resource


@job(resource_defs={
       "wandb_config": make_values_resource(
           entity=str,
           project=str,
       ),
       "wandb_resource": wandb_resource.configured(
           {"api_key": {"env": "WANDB_API_KEY"}}
       ),
   },
)
def run_launch_job_example():
   run_launch_job.alias("my_launched_job")() # rename the job with an alias

Best practices

The following recommendations help you get the most out of the integration once you have it working end to end.

Use the IO Manager to read and write Artifacts

Avoid using Artifact.download() or Run.log_artifact() directly. The integration handles those methods. Instead, return the data you want to store in the Artifact and let the integration do the rest. This approach provides better lineage for the Artifact.

Only build an Artifact object yourself for complex use cases

Return Python objects and W&B objects from your ops and assets. The integration handles bundling the Artifact. For complex use cases, you can build an Artifact directly in a Dagster job. Pass an Artifact object to the integration for metadata enrichment, such as the source integration name and version, the Python version used, the pickle protocol version, and more.

Add files, directories, and external references to your Artifacts through the metadata

Use the integration wandb_artifact_configuration object to add any file, directory, or external references (Amazon S3, GCS, HTTP, and so on). See the advanced example in the Artifact configuration section for more information.

Use an @asset instead of an @op when an Artifact is produced

Artifacts are assets. Use an asset when Dagster maintains that asset. This provides better observability in the Dagit Asset Catalog.

Use a SourceAsset to consume an Artifact created outside Dagster

This lets you take advantage of the integration to read externally created Artifacts. Otherwise, you can only use Artifacts created by the integration.

Use W&B Launch to orchestrate training on dedicated compute for large models

You can train small models inside your Dagster cluster, and you can run Dagster in a Kubernetes cluster with GPU nodes. Use W&B Launch for large model training. This prevents overloading your instance and provides access to more appropriate compute.

Set your W&B Run ID to the value of your Dagster Run ID when experiment tracking within Dagster

Make the Run resumable and set the W&B Run ID to the Dagster Run ID or to a string of your choice. Following this recommendation ensures your W&B metrics and W&B Artifacts are stored in the same W&B Run when you train models inside Dagster.

Either set the W&B Run ID to the Dagster Run ID:

wandb.init(
    id=context.run_id,
    resume="allow",
    ...
)

Or choose your own W&B Run ID and pass it to the IO Manager configuration:

wandb.init(
    id="my_resumable_run_id",
    resume="allow",
    ...
)

@job(
   resource_defs={
       "io_manager": wandb_artifacts_io_manager.configured(
           {"wandb_run_id": "my_resumable_run_id"}
       ),
   }
)

Only collect data you need with get or get_path for large W&B Artifacts

By default, the integration downloads an entire Artifact. If you’re using large artifacts, you might want to collect only the specific files or objects you need. This improves speed and resource utilization.

For Python objects, adapt the pickling module to your use case

By default, the W&B integration uses the standard pickle module. But some objects aren’t compatible with it. For example, functions with yield raise an error if you try to pickle them. W&B supports other Pickle-based serialization modules (dill, cloudpickle, joblib).

You can also use more advanced serialization like ONNX or PMML by returning a serialized string or creating an Artifact directly. The right choice depends on your use case. Refer to the available literature on this subject.