Skip to content

AzureMLLocalRunHook

kedro_azureml_pipeline.hooks.AzureMLLocalRunHook

Hook that configures Azure ML asset datasets for local and remote runs.

See Also

AzureMLAssetDataset : Dataset managed by this hook. AzurePipelinesRunner : Remote runner detected by the hook. WorkspacesConfig : Workspace config injected into datasets.

Source Code

View on GitHub

Show/Hide source 
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
class AzureMLLocalRunHook:
    """Hook that configures Azure ML asset datasets for local and remote runs.

    See Also
    --------
    [AzureMLAssetDataset][kedro_azureml_pipeline.datasets.AzureMLAssetDataset] : Dataset managed by this hook.
    [AzurePipelinesRunner][kedro_azureml_pipeline.runner.AzurePipelinesRunner] : Remote runner detected by the hook.
    [WorkspacesConfig][kedro_azureml_pipeline.config.WorkspacesConfig] : Workspace config injected into datasets.
    """

    @staticmethod
    def _patch_azureml_artifact_builder() -> None:
        """Re-register the ``azureml`` artifact builder with a ``**kwargs``-tolerant wrapper.

        MLflow 3.10+ passes ``tracking_uri`` and ``registry_uri`` keyword arguments
        to artifact repository builders, but ``azureml-mlflow``'s
        ``azureml_artifacts_builder`` does not accept them, causing a ``TypeError``.

        This method wraps the original builder so that extra keyword arguments are
        silently dropped.  The patch is a no-op when:

        * ``mlflow`` or ``azureml-mlflow`` is not installed.
        * The builder already accepts ``**kwargs`` (i.e. azureml-mlflow was updated).
        """
        try:
            from azureml.mlflow.entry_point_loaders import azureml_artifacts_builder
            from mlflow.store.artifact.artifact_repository_registry import _artifact_repository_registry
        except ImportError:
            return

        params = inspect.signature(azureml_artifacts_builder).parameters
        if any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()):
            return

        original = azureml_artifacts_builder

        def _tolerant_builder(artifact_uri=None, **kwargs):  # noqa: ARG001
            """Wrap the original builder, forwarding only ``artifact_uri``."""
            return original(artifact_uri=artifact_uri)

        _artifact_repository_registry.register("azureml", _tolerant_builder)
        logger.info("kedro-azureml-pipeline: patched azureml artifact builder to accept extra kwargs")

    @hook_impl
    def after_context_created(self, context) -> None:
        """Register the ``azureml`` config pattern and resolve workspace config.

        Parameters
        ----------
        context : KedroContext
            Kedro project context.
        """
        self._patch_azureml_artifact_builder()
        if "azureml" not in context.config_loader.config_patterns:
            context.config_loader.config_patterns.update({"azureml": ["azureml*", "azureml*/**", "**/azureml*"]})
        try:
            self.azure_config = WorkspacesConfig.model_validate(context.config_loader["azureml"]["workspace"]).resolve()
        except (KeyError, MissingConfigException):
            self.azure_config = None
            logger.info("kedro-azureml-pipeline: no azureml config found, skipping workspace resolution")

    @hook_impl
    def after_catalog_created(self, catalog):
        """Inject workspace config into all ``AzureMLAssetDataset`` entries.

        Parameters
        ----------
        catalog : DataCatalog
            Created data catalog.
        """
        if self.azure_config is None:
            return
        for dataset_name in catalog.filter():
            dataset = catalog[dataset_name]
            if isinstance(dataset, AzureMLAssetDataset):
                dataset.azure_config = self.azure_config
                catalog[dataset_name] = dataset

    @hook_impl
    def before_pipeline_run(self, run_params, pipeline, catalog):
        """Switch asset datasets between local-intermediate and remote mode.

        Parameters
        ----------
        run_params : dict
            Parameters passed to the run command.
        pipeline : Pipeline
            Pipeline about to be run.
        catalog : DataCatalog
            Data catalog.
        """
        for dataset_name in catalog.filter():
            dataset = catalog[dataset_name]
            if isinstance(dataset, AzureMLAssetDataset):
                if AzurePipelinesRunner.__name__ not in run_params["runner"]:
                    # when running locally using an AzureMLAssetDataset
                    # as an intermediate dataset we don't want download
                    # but still set to run local with a local version.
                    if dataset_name not in pipeline.inputs():
                        dataset.as_local_intermediate()
                # when running remotely we still want to provide information
                # from the azureml config for getting the dataset version during
                # remote runs
                else:
                    dataset.as_remote()

                catalog[dataset_name] = dataset

Methods

after_context_created(context)

Register the azureml config pattern and resolve workspace config.

Parameters
Name Type Description Default
context KedroContext

Kedro project context.

 required
Source Code

View on GitHub

Show/Hide source 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
@hook_impl
def after_context_created(self, context) -> None:
    """Register the ``azureml`` config pattern and resolve workspace config.

    Parameters
    ----------
    context : KedroContext
        Kedro project context.
    """
    self._patch_azureml_artifact_builder()
    if "azureml" not in context.config_loader.config_patterns:
        context.config_loader.config_patterns.update({"azureml": ["azureml*", "azureml*/**", "**/azureml*"]})
    try:
        self.azure_config = WorkspacesConfig.model_validate(context.config_loader["azureml"]["workspace"]).resolve()
    except (KeyError, MissingConfigException):
        self.azure_config = None
        logger.info("kedro-azureml-pipeline: no azureml config found, skipping workspace resolution")

after_catalog_created(catalog)

Inject workspace config into all AzureMLAssetDataset entries.

Parameters
Name Type Description Default
catalog DataCatalog

Created data catalog.

 required
Source Code

View on GitHub

Show/Hide source 
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
@hook_impl
def after_catalog_created(self, catalog):
    """Inject workspace config into all ``AzureMLAssetDataset`` entries.

    Parameters
    ----------
    catalog : DataCatalog
        Created data catalog.
    """
    if self.azure_config is None:
        return
    for dataset_name in catalog.filter():
        dataset = catalog[dataset_name]
        if isinstance(dataset, AzureMLAssetDataset):
            dataset.azure_config = self.azure_config
            catalog[dataset_name] = dataset

before_pipeline_run(run_params, pipeline, catalog)

Switch asset datasets between local-intermediate and remote mode.

Parameters
Name Type Description Default
run_params dict

Parameters passed to the run command.

 required
pipeline Pipeline

Pipeline about to be run.

 required
catalog DataCatalog

Data catalog.

 required
Source Code

View on GitHub

Show/Hide source 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
@hook_impl
def before_pipeline_run(self, run_params, pipeline, catalog):
    """Switch asset datasets between local-intermediate and remote mode.

    Parameters
    ----------
    run_params : dict
        Parameters passed to the run command.
    pipeline : Pipeline
        Pipeline about to be run.
    catalog : DataCatalog
        Data catalog.
    """
    for dataset_name in catalog.filter():
        dataset = catalog[dataset_name]
        if isinstance(dataset, AzureMLAssetDataset):
            if AzurePipelinesRunner.__name__ not in run_params["runner"]:
                # when running locally using an AzureMLAssetDataset
                # as an intermediate dataset we don't want download
                # but still set to run local with a local version.
                if dataset_name not in pipeline.inputs():
                    dataset.as_local_intermediate()
            # when running remotely we still want to provide information
            # from the azureml config for getting the dataset version during
            # remote runs
            else:
                dataset.as_remote()

            catalog[dataset_name] = dataset