Kmonte/setup gs e2e trainers

deployment/configs/e2e_glt_gs_resource_config.yaml

@@ -1,5 +1,6 @@
 # Diffs from e2e_glt_resource_config.yaml
 # - Swap vertex_ai_inferencer_config for vertex_ai_graph_store_inferencer_config
+# - Swap vertex_ai_trainer_config for vertex_ai_graph_store_trainer_config
 shared_resource_config:
   resource_labels:
     cost_resource_group_tag: dev_experiments_COMPONENT
@@ -26,11 +27,17 @@ preprocessor_config:
     machine_type: "n2d-highmem-64"
     disk_size_gb: 300
 trainer_resource_config:
-  vertex_ai_trainer_config:
-    machine_type: n1-highmem-32
-    gpu_type: NVIDIA_TESLA_T4
-    gpu_limit: 2
-    num_replicas: 2
+  vertex_ai_graph_store_trainer_config:
+    graph_store_pool:
+      machine_type: n2-highmem-32
+      gpu_type: ACCELERATOR_TYPE_UNSPECIFIED
+      gpu_limit: 0
+      num_replicas: 2
+    compute_pool:
+      machine_type: n1-standard-16
+      gpu_type: NVIDIA_TESLA_T4
+      gpu_limit: 2
+      num_replicas: 2
 inferencer_resource_config:
   vertex_ai_graph_store_inferencer_config:
     graph_store_pool:

examples/link_prediction/graph_store/configs/e2e_het_dblp_sup_gs_task_config.yaml

@@ -29,7 +29,6 @@ datasetConfig:
     dataPreprocessorArgs:
       # This argument is specific for the `PassthroughPreprocessorConfigForMockedAssets` preprocessor to indicate which dataset we should be using
       mocked_dataset_name: 'dblp_node_anchor_edge_features_lp'
-# TODO(kmonte): Add GS trainer
 trainerConfig:
   trainerArgs:
     # Example argument to trainer
@@ -49,7 +48,15 @@ trainerConfig:
         ("paper", "to", "author"): [15, 15],
         ("author", "to", "paper"): [20, 20]
       }
-  command: python -m examples.link_prediction.heterogeneous_training
+  command: python -m examples.link_prediction.graph_store.heterogeneous_training
+  graphStoreStorageConfig:
+    command: python -m examples.link_prediction.graph_store.storage_main
+    storageArgs:
+      sample_edge_direction: "in"
+      splitter_cls_path: "gigl.utils.data_splitters.DistNodeAnchorLinkSplitter"
+      splitter_kwargs: '{"sampling_direction": "in", "should_convert_labels_to_edges": True, "num_val": 0.3, "num_test": 0.3, "supervision_edge_types": [("author", "to", "paper")]}'
+      ssl_positive_label_percentage: "0.15"
+      num_server_sessions: "1"
 # TODO(kmonte): Move to user-defined server code
 inferencerConfig:
   inferencerArgs:

examples/link_prediction/graph_store/configs/e2e_hom_cora_sup_gs_task_config.yaml

@@ -15,13 +15,19 @@ datasetConfig:
     dataPreprocessorArgs:
       # This argument is specific for the `PassthroughPreprocessorConfigForMockedAssets` preprocessor to indicate which dataset we should be using
       mocked_dataset_name: 'cora_homogeneous_node_anchor_edge_features_user_defined_labels'
-# TODO(kmonte): Add GS trainer
 trainerConfig:
   trainerArgs:
     # Example argument to trainer
     log_every_n_batch: "50" # Frequency in which we log batch information
     num_neighbors: "[10, 10]" # Fanout per hop, specified as a string representation of a list for the homogeneous use case
-  command: python -m examples.link_prediction.homogeneous_training
+  command: python -m examples.link_prediction.graph_store.homogeneous_training
+  graphStoreStorageConfig:
+    command: python -m examples.link_prediction.graph_store.storage_main
+    storageArgs:
+      sample_edge_direction: "in"
+      splitter_cls_path: "gigl.utils.data_splitters.DistNodeAnchorLinkSplitter"
+      splitter_kwargs: '{"sampling_direction": "in", "should_convert_labels_to_edges": True, "num_val": 0.1, "num_test": 0.1}'
+      num_server_sessions: "1"
 # TODO(kmonte): Move to user-defined server code
 inferencerConfig:
   inferencerArgs:

examples/link_prediction/graph_store/configs/example_resource_config.yaml

@@ -46,13 +46,18 @@ preprocessor_config:
     max_num_workers: 4
     machine_type: "n2-standard-16"
     disk_size_gb: 300
-# TODO(kmonte): Update
 trainer_resource_config:
-  vertex_ai_trainer_config:
-    machine_type: n1-standard-16
-    gpu_type: NVIDIA_TESLA_T4
-    gpu_limit: 2
-    num_replicas: 2
+  vertex_ai_graph_store_trainer_config:
+    graph_store_pool:
+      machine_type: n2-highmem-32
+      gpu_type: ACCELERATOR_TYPE_UNSPECIFIED
+      gpu_limit: 0
+      num_replicas: 2
+    compute_pool:
+      machine_type: n1-standard-16
+      gpu_type: NVIDIA_TESLA_T4
+      gpu_limit: 2
+      num_replicas: 2
 inferencer_resource_config:
   vertex_ai_graph_store_inferencer_config:
     graph_store_pool: