DOCSEARCHbyAlgolia
Submit Docs

hydra-zen Documentation

repository·main·Indexed 19 days ago

https://github.com/mit-ll-responsible-ai/hydra-zen

A Python library that simplifies writing configurable, repeatable, and scalable workflows by building on top of Hydra. It eliminates manual YAML configuration by dynamically generating dataclass-based configs using functions like make_config and builds. Key features include a custom config-store API, a task-function wrapper (hydra_zen.zen), and support for automatic type refinement and runtime data validation via integration with Pydantic and beartype.

Tokens
32.6K
Snippets
94
Records
117
Agent score
64%

What's inside hydra-zen

  1. Overview of hydra-zen

    main

    hydra-zen is a Python library designed to facilitate configurable, repeatable, and scalable workflows by building on top of Hydra.

    Key features include:

    • Configurability: Manage all aspects of your code via a single interface (CLI or Python function).
    • Repeatability: Automatically saves the full configuration alongside results for self-documenting runs.
    • Scalability: Supports launching multiple runs locally or across cluster nodes.

    A core value proposition is the elimination of hand-written YAML configs. Instead, hydra-zen uses functions to dynamically generate dataclass-based configurations, provides a custom config-store API, and offers a task-function wrapper (hydra_zen.zen) to reduce Hydra-specific boilerplate.

  2. The hydra-zen configuration workflow

    main

    The core workflow in hydra-zen follows these steps:

    1. Create Configs: Use functions like make_config and builds to define your configuration structures.
    2. Store Configs: Use hydra_zen.store to register these configs so they can be accessed by Hydra.
    3. Launch Jobs: Use zen or launch to run a Hydra job directly from a Python function.
    4. Instantiate/Resolve: Use instantiate to turn the configuration objects into actual data and class instances.

  3. Keep configurations DRY using `builds`

    main

    Manually defining configurations via YAML files or Python dataclasses is considered WET (Writing Everything Twice) because you must manually mirror the class's import path, parameter names, type annotations, and default values. If the underlying class changes, the configuration becomes out of sync, leading to bugs.

    To stay DRY (Don't Repeat Yourself), use hydra_zen.builds to dynamically generate configurations. This approach automatically inspects the target class at runtime to create a configuration that is always in sync with the class signature, including type annotations and default values. Additionally, builds validates any manually specified parameters against the target class's signature during configuration creation, catching typos and mistakes early.

    from hydra_zen import builds
from vision.model import DNN

# Dynamically generates a configuration for the DNN class
ZenBuilds_DNN = builds(DNN, populate_full_signature=True)

  4. How hydra-zen works: Core Concepts

    main

    hydra-zen is designed to eliminate hand-written YAML configurations by providing a Python-native way to manage Hydra configs.

    Key Abstractions:

    • Configurability: Uses ZenStore to manage hierarchical, nested parameters that can be swapped via CLI.
    • Repeatability: Automatically saves the full configuration used for a run as a YAML file in a Hydra-managed output directory.
    • ZenStore: A custom config-store API that allows you to define and group configurations in Python code.
    • zen() wrapper: A task-function wrapper that simplifies the creation of the Hydra entry point and CLI.
    • just(): A utility to specify that a value (like a function or class) should be treated as a reference/import rather than being instantiated with parameters.

  5. Generate hierarchical configs with builds()

    main

    You can use hydra_zen.builds to dynamically generate configuration objects that mirror the interface of existing Python classes or functions. This allows your configuration structure to stay in sync with your library code.

    Key features:

    • populate_full_signature=True: When passed to builds, this ensures the generated config includes all arguments from the target's signature (including defaults).
    • Nesting: You can nest one generated config inside another by passing a config object as an argument to a subsequent builds call.
    from hydra_zen import builds
from game_library import Character, inventory

# Create a builder for the inventory function
InventoryConf = builds(inventory, populate_full_signature=True)
starter_gear = InventoryConf(gold=10, weapon="stick", costume="tunic")

# Create a builder for the Character class, nesting the inventory config
CharConf = builds(Character, inventory=starter_gear, populate_full_signature=True)

  6. Inspect Hydra application outputs and working directory

    main

    When you launch an application, Hydra creates an outputs/ directory containing time-stamped subdirectories. Each job's specific directory is accessible via job.working_dir.

    Inside the job directory, you will find:

    • Your application's generated files (e.g., player_log.txt).
    • A .hydra/ directory containing YAML files for the specific run:
      • config.yaml: The final configuration used.
      • hydra.yaml: Hydra's internal configuration.
      • overrides.yaml: The overrides applied during launch.
    # Access the directory where the job ran
job_dir = Path(job.working_dir)

# List contents (includes .hydra and application logs)
print(sorted(job_dir.glob("*")))

# Inspect the final configuration used for this run
# print_file is a helper for this example
print_file(job_dir / ".hydra" / "config.yaml")

  7. How `zen_wrappers` are documented in Hydra configs

    main

    When using zen_wrappers via hydra_zen.builds, the resulting Hydra configuration file (config.yaml) explicitly documents the injection. The configuration will include:

    • _target_: The internal hydra-zen processing target.
    • _zen_target: The original class being instantiated.
    • _zen_wrappers: The name of the wrapper function applied.

    This ensures that any job run using the generated configuration is self-documenting and reproducible, as the exact transformations applied to the objects are recorded in the job's metadata.

  8. Use @hydrated_dataclass to combine static and dynamic configurations

    main

    The @hydrated_dataclass decorator allows you to create configuration classes that combine the benefits of Python dataclasses with the dynamic auto-population capabilities of hydra-zen.builds.

    Key benefits include:

    • Static Analysis Support: Attributes are statically available to IDEs and type-checkers (like pyright), allowing them to catch type mismatches or mutations of frozen=True classes.
    • Automatic Hydra Metadata: The decorator automatically handles Hydra-specific fields like _target_ based on the target argument.
    • Runtime Validation: It provides runtime validation upon construction, catching misspelled parameter names that do not match the target class's signature.

    To use it, decorate a class with @hydrated_dataclass(target=...). You can pass arguments like zen_partial=True to create a partial function or frozen=True to make the configuration immutable.

    from hydra_zen import hydrated_dataclass
from torch.optim import Adam

@hydrated_dataclass(target=Adam, zen_partial=True, frozen=True)
class BuildsAdam:
    lr: float = 0.01
    momentum: float = 0.9

# This creates a partial function for Adam with the specified parameters
from hydra_zen import instantiate
partial_adam = instantiate(BuildsAdam)
# Result: functools.partial(<class Adam>, lr=0.01, momentum=0.9)

  9. How Hydra's callback system works with hydra-zen

    main

    Hydra's callback system allows you to run custom code triggered by specific events, such as on_job_start or on_job_end. This is useful for modular tasks like performance profiling, uploading results to cloud storage, or logging, which are independent of your main task function.

    In hydra-zen, you can integrate these callbacks by:

    1. Defining a class that inherits from hydra.experimental.callback.Callback.
    2. Registering these callbacks in a ZenStore.
    3. Either adding them directly to the HydraConf (to enable them by default) or adding them to a specific configuration group (to enable them via the CLI).

    Callbacks can access the job's configuration (e.g., the Config dataclass of your task) during lifecycle events like on_job_end.

    from hydra.experimental.callback import Callback

class MyCallback(Callback):
    def on_job_start(self, **kw) -> None:
        # logic before job starts
        pass

    def on_job_end(self, config: MyTaskConfig, **kwargs) -> None:
        # logic after job ends, with access to task config
        pass

  10. How hydra-zen handles Hydra's limited type support

    main

    Hydra has a narrow subset of supported type annotations (e.g., Any, primitives like int/str, Enums, List, Dict, and Optional). Using complex annotations like Literal directly in a dataclass or function signature will cause Hydra to raise a ConfigTypeError during instantiation.

    To solve this, hydra-zen's config-creation functions (like builds) perform Automatic Type Refinement. They automatically broaden unsupported type annotations to compatible ones (e.g., broadening List[Literal[1, 2]] to List[Any]) so that Hydra can process the config without errors, while still preserving as much type information as possible for validation.

    from typing_extensions import Literal
from dataclasses import dataclass

@dataclass
class A:
    x: Literal[1, 2]

# This would normally fail in Hydra:
# instantiate(A, x=1) -> ConfigTypeError

  11. Use auto-config with `store()`

    main

    Hydra-zen's store() function has auto-config capabilities that allow you to create configurations more concisely. When using auto-config, you can pass arguments directly to store() to populate the configuration. This is equivalent to manually wrapping the function with builds(..., populate_full_signature=True) before storing it.

    With auto-config:

    from hydra_zen import store

def func(x, y):
   ...

store(func, x=2, y=3)

    Without auto-config (explicit):

    from hydra_zen import builds, store

def func(x, y):
   ...

store(builds(func, x=2, y=3, populate_full_signature=True), name="func")
    from hydra_zen import store

def func(x, y):
   ...

store(func, x=2, y=3)

  12. Create configs with make_config and builds

    main

    Use make_config and builds to create Hydra-compatible configuration objects. These functions allow you to use complex Python types that Hydra cannot natively serialize by automatically converting them into Hydra-compatible dataclass representations.

    Key functions:

    • make_config: Creates a configuration object.
    • builds: Creates a configuration object, often used for building complex structures or partial functions.
    from hydra_zen import builds, just, make_config, to_yaml, instantiate

# Example: complex numbers
Conf = make_config(value=2.0 + 3.0j)

# Example: partial functions
from functools import partial
Conf2 = builds(dict, x=partial(int, 3))

# Example: dataclasses
from dataclasses import dataclass
@dataclass
class Bar:
    reduce_fn: callable = sum
just_bar = just(Bar())