Move graphsage endpoints under shared surface

Author: FlorentinD

graphdatascience/procedure_surface/api/node_embedding/graphsage_endpoints.py

@@ -4,19 +4,17 @@
 
 from graphdatascience.procedure_surface.api.catalog.graph_api import GraphV2
 from graphdatascience.procedure_surface.api.estimation_result import EstimationResult
-from graphdatascience.procedure_surface.api.model.graphsage_model import GraphSageModelV2
 from graphdatascience.procedure_surface.api.node_embedding.graphsage_predict_endpoints import (
     GraphSageMutateResult,
     GraphSagePredictEndpoints,
     GraphSageWriteResult,
 )
 from graphdatascience.procedure_surface.api.node_embedding.graphsage_train_endpoints import (
     GraphSageTrainEndpoints,
-    GraphSageTrainResult,
 )
 
 
-class GraphSageEndpoints(GraphSageTrainEndpoints, GraphSagePredictEndpoints):
+class GraphSageEndpoints(GraphSagePredictEndpoints):
     """
     API for the GraphSage algorithm, combining both training and prediction functionalities.
     """
@@ -29,66 +27,74 @@ def __init__(
         self._train_endpoints = train_endpoints
         self._predict_endpoints = predict_endpoints
 
-    def train(
-        self,
-        G: GraphV2,
-        model_name: str,
-        feature_properties: list[str],
-        *,
-        activation_function: Any | None = None,
-        negative_sample_weight: int | None = None,
-        embedding_dimension: int | None = None,
-        tolerance: float | None = None,
-        learning_rate: float | None = None,
-        max_iterations: int | None = None,
-        sample_sizes: list[int] | None = None,
-        aggregator: Any | None = None,
-        penalty_l2: float | None = None,
-        search_depth: int | None = None,
-        epochs: int | None = None,
-        projected_feature_dimension: int | None = None,
-        batch_sampling_ratio: float | None = None,
-        store_model_to_disk: bool | None = None,
-        relationship_types: list[str] | None = None,
-        node_labels: list[str] | None = None,
-        username: str | None = None,
-        log_progress: bool = True,
-        sudo: bool | None = None,
-        concurrency: Any | None = None,
-        job_id: Any | None = None,
-        batch_size: int | None = None,
-        relationship_weight_property: str | None = None,
-        random_seed: Any | None = None,
-    ) -> tuple[GraphSageModelV2, GraphSageTrainResult]:
-        return self._train_endpoints.train(
-            G,
-            model_name,
-            feature_properties,
-            activation_function=activation_function,
-            negative_sample_weight=negative_sample_weight,
-            embedding_dimension=embedding_dimension,
-            tolerance=tolerance,
-            learning_rate=learning_rate,
-            max_iterations=max_iterations,
-            sample_sizes=sample_sizes,
-            aggregator=aggregator,
-            penalty_l2=penalty_l2,
-            search_depth=search_depth,
-            epochs=epochs,
-            projected_feature_dimension=projected_feature_dimension,
-            batch_sampling_ratio=batch_sampling_ratio,
-            store_model_to_disk=store_model_to_disk,
-            relationship_types=relationship_types,
-            node_labels=node_labels,
-            username=username,
-            log_progress=log_progress,
-            sudo=sudo,
-            concurrency=concurrency,
-            job_id=job_id,
-            batch_size=batch_size,
-            relationship_weight_property=relationship_weight_property,
-            random_seed=random_seed,
-        )
+    @property
+    def train(self) -> GraphSageTrainEndpoints:
+        """
+        Trains a GraphSage model on the given graph.
+
+        Parameters
+        ----------
+        G : GraphV2
+            The graph to run the algorithm on
+        model_name : str
+            Name under which the model will be stored
+        feature_properties : list[str]
+            The names of the node properties to use as input features
+        activation_function : Any | None, default=None
+            The activation function to apply after each layer
+        negative_sample_weight : int | None, default=None
+            Weight of negative samples in the loss function
+        embedding_dimension : int | None, default=None
+            The dimension of the generated embeddings
+        tolerance : float | None, default=None
+            Tolerance for early stopping based on loss improvement
+        learning_rate : float | None, default=None
+            Learning rate for the training optimization
+        max_iterations : int | None, default=None
+            Maximum number of training iterations
+        sample_sizes : list[int] | None, default=None
+            Number of neighbors to sample at each layer
+        aggregator : Any | None, default=None
+            The aggregator function for neighborhood aggregation
+        penalty_l2 : float | None, default=None
+            L2 regularization penalty
+        search_depth : int | None, default=None
+            Maximum search depth for neighbor sampling
+        epochs : int | None, default=None
+            Number of training epochs
+        projected_feature_dimension : int | None, default=None
+            Dimension to project input features to before training
+        batch_sampling_ratio : float | None, default=None
+            Ratio of nodes to sample for each training batch
+        store_model_to_disk : bool | None, default=None
+            Whether to persist the model to disk
+        relationship_types : list[str] | None, default=None
+            The relationship types used to select relationships for this algorithm run
+        node_labels : list[str] | None, default=None
+            The node labels used to select nodes for this algorithm run
+        username : str | None = None
+            The username to attribute the procedure run to
+        log_progress : bool | None, default=None
+            Whether to log progress
+        sudo : bool | None, default=None
+            Override memory estimation limits
+        concurrency : Any | None, default=None
+            The number of concurrent threads
+        job_id : Any | None, default=None
+            An identifier for the job
+        batch_size : int | None, default=None
+            Batch size for training
+        relationship_weight_property : str | None, default=None
+            The property name that contains weight
+        random_seed : Any | None, default=None
+            Random seed for reproducible results
+
+        Returns
+        -------
+        GraphSageModelV2
+            Trained model
+        """
+        return self._train_endpoints
 
     def stream(
         self,

graphdatascience/procedure_surface/api/node_embedding/graphsage_train_endpoints.py

@@ -5,16 +5,46 @@
 
 from graphdatascience.procedure_surface.api.base_result import BaseResult
 from graphdatascience.procedure_surface.api.catalog.graph_api import GraphV2
+from graphdatascience.procedure_surface.api.estimation_result import EstimationResult
 from graphdatascience.procedure_surface.api.model.graphsage_model import GraphSageModelV2
 
 
 class GraphSageTrainEndpoints(ABC):
-    """
-    Abstract base class defining the API for the GraphSage algorithm.
-    """
+    @abstractmethod
+    def __call__(
+        self,
+        G: GraphV2,
+        model_name: str,
+        feature_properties: list[str],
+        *,
+        activation_function: Any | None = None,
+        negative_sample_weight: int | None = None,
+        embedding_dimension: int | None = None,
+        tolerance: float | None = None,
+        learning_rate: float | None = None,
+        max_iterations: int | None = None,
+        sample_sizes: list[int] | None = None,
+        aggregator: Any | None = None,
+        penalty_l2: float | None = None,
+        search_depth: int | None = None,
+        epochs: int | None = None,
+        projected_feature_dimension: int | None = None,
+        batch_sampling_ratio: float | None = None,
+        store_model_to_disk: bool | None = None,
+        relationship_types: list[str] | None = None,
+        node_labels: list[str] | None = None,
+        username: str | None = None,
+        log_progress: bool = True,
+        sudo: bool | None = None,
+        concurrency: Any | None = None,
+        job_id: Any | None = None,
+        batch_size: int | None = None,
+        relationship_weight_property: str | None = None,
+        random_seed: Any | None = None,
+    ) -> tuple[GraphSageModelV2, GraphSageTrainResult]: ...
 
     @abstractmethod
-    def train(
+    def estimate(
         self,
         G: GraphV2,
         model_name: str,
@@ -44,9 +74,13 @@ def train(
         batch_size: int | None = None,
         relationship_weight_property: str | None = None,
         random_seed: Any | None = None,
-    ) -> tuple[GraphSageModelV2, GraphSageTrainResult]:
+    ) -> EstimationResult:
         """
-        Trains a GraphSage model on the given graph.
+        Estimates memory requirements and other statistics for training a GraphSage model.
+
+        This method provides memory estimation for the GraphSage training algorithm without
+        actually executing the training. It helps determine the computational requirements
+        before running the actual training procedure.
 
         Parameters
         ----------
@@ -84,9 +118,9 @@ def train(
             Ratio of nodes to sample for each training batch
         store_model_to_disk : bool | None, default=None
             Whether to persist the model to disk
-        relationship_types : list[str] | None, default=None
+        relationship_types : list[str] | None = None
             The relationship types used to select relationships for this algorithm run
-        node_labels : list[str] | None, default=None
+        node_labels : list[str] | None = None
             The node labels used to select nodes for this algorithm run
         username : str | None = None
             The username to attribute the procedure run to
@@ -107,8 +141,8 @@ def train(
 
         Returns
         -------
-        GraphSageModelV2
-            Trained model
+        EstimationResult
+            The estimation result containing memory requirements and other statistics
         """
 
 

graphdatascience/procedure_surface/arrow/node_embedding/graphsage_train_arrow_endpoints.py

@@ -3,6 +3,7 @@
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
 from graphdatascience.arrow_client.v2.remote_write_back_client import RemoteWriteBackClient
 from graphdatascience.procedure_surface.api.catalog.graph_api import GraphV2
+from graphdatascience.procedure_surface.api.estimation_result import EstimationResult
 from graphdatascience.procedure_surface.api.model.graphsage_model import GraphSageModelV2
 from graphdatascience.procedure_surface.api.node_embedding.graphsage_train_endpoints import (
     GraphSageTrainEndpoints,
@@ -29,7 +30,7 @@ def __init__(
         )
         self._model_api = ModelApiArrow(arrow_client)
 
-    def train(
+    def __call__(
         self,
         G: GraphV2,
         model_name: str,
@@ -100,3 +101,67 @@ def train(
         train_result = GraphSageTrainResult(**result)
 
         return model, train_result
+
+    def estimate(
+        self,
+        G: GraphV2,
+        model_name: str,
+        feature_properties: list[str],
+        *,
+        activation_function: Any | None = None,
+        negative_sample_weight: int | None = None,
+        embedding_dimension: int | None = None,
+        tolerance: float | None = None,
+        learning_rate: float | None = None,
+        max_iterations: int | None = None,
+        sample_sizes: list[int] | None = None,
+        aggregator: Any | None = None,
+        penalty_l2: float | None = None,
+        search_depth: int | None = None,
+        epochs: int | None = None,
+        projected_feature_dimension: int | None = None,
+        batch_sampling_ratio: float | None = None,
+        store_model_to_disk: bool | None = None,
+        relationship_types: list[str] | None = None,
+        node_labels: list[str] | None = None,
+        username: str | None = None,
+        log_progress: bool = True,
+        sudo: bool | None = None,
+        concurrency: Any | None = None,
+        job_id: Any | None = None,
+        batch_size: int | None = None,
+        relationship_weight_property: str | None = None,
+        random_seed: Any | None = None,
+    ) -> EstimationResult:
+        return self._node_property_endpoints.estimate(
+            estimate_endpoint="v2/embeddings.graphSage.train.estimate",
+            G=G,
+            algo_config=self._node_property_endpoints.create_estimate_config(
+                model_name=model_name,
+                feature_properties=feature_properties,
+                activation_function=activation_function,
+                negative_sample_weight=negative_sample_weight,
+                embedding_dimension=embedding_dimension,
+                tolerance=tolerance,
+                learning_rate=learning_rate,
+                max_iterations=max_iterations,
+                sample_sizes=sample_sizes,
+                aggregator=aggregator,
+                penalty_l2=penalty_l2,
+                search_depth=search_depth,
+                epochs=epochs,
+                projected_feature_dimension=projected_feature_dimension,
+                batch_sampling_ratio=batch_sampling_ratio,
+                store_model_to_disk=store_model_to_disk,
+                relationship_types=relationship_types,
+                node_labels=node_labels,
+                username=username,
+                log_progress=log_progress,
+                sudo=sudo,
+                concurrency=concurrency,
+                job_id=job_id,
+                batch_size=batch_size,
+                relationship_weight_property=relationship_weight_property,
+                random_seed=random_seed,
+            ),
+        )