cursus.steps.configs.config_graphstorm_gnn_tuning_step

GraphStorm GNN Tuning Configuration (the Tuning verb’s first interface, FZ 31e1d3p / plan 31e1d3p2).

A hyperparameter-tuning step over GraphStorm GNN training. A tuning job is a training job with a search wrapper, so this config is the idiomatic composition of the wrapped training config (GraphStormGNNTrainingConfig — every estimator field the byo_container factory reads) and the search-axis mixin (TuningStepConfigMixin — objective / search_space / strategy / limits / metric_definitions). TuningHandler reads both off b.config at build time: it builds the SAME GraphStorm estimator GraphStormGNNTraining would, then wraps it in a HyperparameterTuner.

Because the estimator is a byo_container (a custom GraphStorm image, no SDK-inferred metrics), metric_definitions is mandatory — the objective is regex-scraped from the container’s stdout, exactly as the Nexus launch_hpo.py launcher does.

class GraphStormGNNTuningConfig(*, objective_metric_name, search_space, objective_type='Maximize', tuning_strategy='Bayesian', max_jobs=20, max_parallel_jobs=1, early_stopping_type='Off', metric_definitions=None, author, bucket, role, region, service_name, pipeline_version, model_class='xgboost', current_date=<factory>, framework_version='2.1.0', py_version='py310', image_uri=None, subnets=None, security_group_ids=None, enable_network_isolation=None, source_dir=None, enable_caching=False, use_secure_pypi=False, max_runtime_seconds=172800, project_root_folder, training_entry_point, training_image_uri, training_instance_type='ml.g5.12xlarge', training_instance_count=1, training_volume_size=125, training_mode='multi_task', num_servers=1, code_s3_uri='', batch_size=1024, hidden_size=100, num_layers=3, fanout='30, 30, 30', **extra_data)[source]

Bases: GraphStormGNNTrainingConfig, TuningStepConfigMixin

GraphStorm/DGL R-GCN hyperparameter tuning in a bring-your-own GraphStorm container.

Inherits every training field (training_image_uri, instance type/count, training_mode, num_servers …) from GraphStormGNNTrainingConfig and adds the search-axis fields (objective_metric_name, search_space, tuning_strategy, max_jobs, max_parallel_jobs, early_stopping_type, metric_definitions) from TuningStepConfigMixin. The mixin’s validator enforces that metric_definitions is present (byo_container detected via the inherited training_image_uri).

model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True, 'extra': 'allow', 'protected_namespaces': (), 'validate_assignment': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_post_init(context, /)

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

  • self (BaseModel) – The BaseModel instance.

  • context (Any) – The context.