cursus.steps.configs.config_graphstorm_gnn_training_step

GraphStorm GNN Training Configuration with Self-Contained Derivation Logic.

Configuration for the GraphStormGNNTraining Training step — GraphStorm/DGL R-GCN node-classification (or multi-task) training over a partitioned DGL heterograph, run in a bring-your-own GraphStorm ECR container. Extends BasePipelineConfig (as PyTorchTrainingConfig does).

The crux is training_image_uri — the byo_container builder reads it (interface image_uri_field: training_image_uri) and sets it VERBATIM as AlgorithmSpecification.TrainingImage, with no image_uris.retrieve (the GraphStorm/DGL stack is in no AWS DLC). graphstorm/dgl/torch are baked into the image, so there is no framework install here.

class GraphStormGNNTrainingConfig(*, 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: BasePipelineConfig

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

training_entry_point: str
training_image_uri: str
training_instance_type: str
training_instance_count: int
training_volume_size: int
max_runtime_seconds: int
training_mode: str
num_servers: int
code_s3_uri: str
batch_size: int
hidden_size: int
num_layers: int
fanout: str
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].

property batch_size_override: int | None

Pre-adjust batch_size for the instance’s GPU memory (ports launch_training._adjust_batch_size). Returns the reduced batch size only when it is smaller than the configured batch_size; else None (emitted as the BATCH_SIZE_OVERRIDE env var only when set).

get_environment_variables()[source]

Env vars the container reads (TRAINING_MODE / NUM_SERVERS / optional BATCH_SIZE_OVERRIDE).

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.