> ## Documentation Index > Fetch the complete documentation index at: https://wb-21fd5541-docs-2626.mintlify.site/llms.txt > Use this file to discover all available pages before exploring further. > Integrate W&B with Dagster to track ML experiments and manage data pipelines with automatic logging and monitoring. # Dagster Use Dagster and W\&B together to orchestrate MLOps pipelines and maintain ML assets. Within your existing Dagster workflows, you can track experiments, manage data, and run training jobs. The W\&B integration supports the following tasks: * Create and use a [W\&B Artifact](/models/artifacts). * Use and create Registered Models in [W\&B Registry](/models/registry). * Run training jobs on dedicated compute using [W\&B Launch](/platform/launch). * Use the [wandb](/models/ref/python) 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. You'll learn how to set up Dagster, create and use W\&B Artifacts in ops and assets, use W\&B Launch, and apply best practices. ## Before you get started You need the following resources to use Dagster within W\&B: * **W\&B API Key**. * **W\&B entity**: The name of your [W\&B Team](/platform/hosting/iam/org_team_struct). * **W\&B project**: The name of the project where [W\&B Runs](/models/runs) are stored. Find your W\&B entity by checking your team's profile page in W\&B. You can use a pre-existing W\&B project or create a new one. Create new projects on the W\&B homepage or on your team's 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 W\&B using an API key, which you must make available to Dagster as an environment variable. 1. [Log in to W\&B](https://wandb.ai/login). If you're using W\&B Server, ask your admin for the instance host name. 2. Create an API key at [User Settings](https://wandb.ai/settings). For a production environment, use a [service account](/support/models/articles/what-is-a-service-account-and-why-is-it-) to own that key. 3. Set an environment variable for that API key: `export WANDB_API_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 [Configuration](#configuration) section. Example configuration for `@job`: ```python theme={null} # 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() ``` Example configuration for `@repository` using assets: ```python theme={null} 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](https://docs.dagster.io/guides/build/external-resources) 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](https://docs.dagster.io/guides/build/io-managers) 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](https://docs.dagster.io/guides/build/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. ```python theme={null} @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](/models/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`): Anything that can be serialized with the [pickle](https://docs.python.org/3/library/pickle.html) 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](#read-wb-artifacts) for more details). ```python theme={null} @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](https://docs.python.org/3/library/pickle.html), [dill](https://github.com/uqfoundation/dill), [cloudpickle](https://github.com/cloudpipe/cloudpickle), [joblib](https://github.com/joblib/joblib)). You can also use more advanced serialization like [ONNX](https://onnx.ai/) or [PMML](https://en.wikipedia.org/wiki/Predictive_Model_Markup_Language). Refer to the [Serialization](#serialization-configuration) section for more information. Any W\&B object, such as a [Table](/models/ref/python/data-types/table) or [Image](/models/ref/python/data-types/image), is added to an Artifact created by the integration. This example adds a Table to an Artifact: ```python theme={null} 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]]) ``` 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. ```python theme={null} 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](https://docs.dagster.io/_apidocs/ops#dagster.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](https://docs.dagster.io/_apidocs/assets#dagster.AssetOut) metadata arguments. The following code examples demonstrate how to configure a dictionary on an `@op`, `@asset`, and `@multi_asset` computations: Example for `@op`: ```python theme={null} @op( out=Out( metadata={ "wandb_artifact_configuration": { "name": "my_artifact", "type": "dataset", } } ) ) def create_dataset(): return [1, 2, 3] ``` Example for `@asset`: ```python theme={null} @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`: ```python theme={null} @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`](/models/ref/python/experiments/artifact#add_dir): (list\[dict\[str, Any]]): An array containing configuration for each local directory to include in the Artifact. * [`add_files`](/models/ref/python/experiments/artifact#add_file): (list\[dict\[str, Any]]): An array containing configuration for each local file to include in the Artifact. * [`add_references`](/models/ref/python/experiments/artifact#add_reference): (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: ```python theme={null} @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. Dagster's UI with an asset details view with attached W&B metadata, including references to a W&B project and run 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. W&B Artifact page with enriched configuration metadata from Dagster W&B Artifact metadata panel with additional configuration details from Dagster W&B Artifact view with further enriched configuration metadata fields from Dagster If you use a static type checker like mypy, import the configuration type definition object using the following: ```python theme={null} from dagster_wandb import WandbArtifactConfiguration ``` ### Use partitions The integration directly supports [Dagster partitions](https://docs.dagster.io/guides/build/partitions-and-backfills). The following is an example with a partitioned using `DailyPartitionsDefinition`: ```python theme={null} @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. Dagster UI with multiple runs for partitioned assets, each partition as a separate run #### Advanced usage For advanced usage, see the following examples: * [Partitioned job](https://github.com/dagster-io/dagster/blob/master/examples/with_wandb/with_wandb/ops/partitioned_job.py) * [Simple partitioned asset](https://github.com/wandb/dagster/blob/master/examples/with_wandb/with_wandb/assets/simple_partitions_example.py) * [Multi-partitioned asset](https://github.com/wandb/dagster/blob/master/examples/with_wandb/with_wandb/assets/multi_partitions_example.py) * [Advanced partitioned usage](https://github.com/wandb/dagster/blob/master/examples/with_wandb/with_wandb/assets/advanced_partitions_example.py) ## 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](https://docs.dagster.io/_apidocs/ops#dagster.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](https://docs.dagster.io/_apidocs/assets#dagster.AssetIn) 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](https://docs.dagster.io/_apidocs/assets#dagster.SourceAsset). It always reads the latest version of that asset. The following examples demonstrate how to read an Artifact from various ops. Reading an artifact from an `@op`: ```python theme={null} @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) ``` Reading an artifact created by another `@asset`: ```python theme={null} @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) ``` Reading an Artifact created outside Dagster: ```python theme={null} 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`: ```python theme={null} @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`: ```python theme={null} @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): ```python theme={null} @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](https://docs.python.org/3/library/pickle.html) 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](https://github.com/uqfoundation/dill), [cloudpickle](https://github.com/cloudpipe/cloudpickle), [joblib](https://github.com/joblib/joblib)). You can also use more advanced serialization like [ONNX](https://onnx.ai/) or [PMML](https://en.wikipedia.org/wiki/Predictive_Model_Markup_Language) 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. ```python theme={null} @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. ```python theme={null} 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](https://docs.dagster.io/guides/build/partitions-and-backfills). 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. 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. ```python theme={null} @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}") ``` The `AssetIn`'s `partition_mapping` configuration lets you choose specific partitions. In this case, the example uses the `TimeWindowPartitionMapping`. ```python theme={null} @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: ```python theme={null} "*": { "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: ```python theme={null} "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: ```python theme={null} "orange": { "version": "v0", }, ``` This configuration fetches data from version `v0` of the `orange` Artifact partition. For aliases: ```python theme={null} "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](https://github.com/dagster-io/dagster/blob/master/examples/with_wandb/with_wandb/assets/advanced_example.py) * [Partitioned job example](https://github.com/dagster-io/dagster/blob/master/examples/with_wandb/with_wandb/ops/partitioned_job.py) * [Linking a model to the Model Registry](https://github.com/dagster-io/dagster/blob/master/examples/with_wandb/with_wandb/assets/model_registry_example.py) ## 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](/platform/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](/platform/launch). You can also view descriptions for all properties in Launchpad. W&B Launchpad interface with agent configuration options and descriptions for Dagster integration Example: ```python theme={null} # 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](/platform/launch). You can also view descriptions for all properties in Launchpad. W&B Launchpad interface with job configuration options and descriptions for Dagster integration Example: ```python theme={null} # 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()`](/models/ref/python/experiments/artifact#download) or [`Run.log_artifact()`](/models/ref/python/experiments/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](#configuration-1) 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](/models/runs/resuming) 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: ```python theme={null} wandb.init( id=context.run_id, resume="allow", ... ) ``` Or choose your own W\&B Run ID and pass it to the IO Manager configuration: ```python theme={null} 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](https://docs.python.org/3/library/pickle.html) 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](https://github.com/uqfoundation/dill), [cloudpickle](https://github.com/cloudpipe/cloudpickle), [joblib](https://github.com/joblib/joblib)). You can also use more advanced serialization like [ONNX](https://onnx.ai/) or [PMML](https://en.wikipedia.org/wiki/Predictive_Model_Markup_Language) 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.