relationalai 0.11.2__py3-none-any.whl → 0.11.4__py3-none-any.whl

relationalai/semantics/reasoners/graph/core.py

@@ -2476,7 +2476,7 @@ class Graph():
 
 
     @include_in_docs
-    def weighted_degree(self):
+    def weighted_degree(self, *, of: Optional[Relationship] = None):
         """Returns a binary relationship containing the weighted degree of each node.
 
         A node's weighted degree is the sum of the weights of all edges
@@ -2484,6 +2484,13 @@ class Graph():
         of both incoming and outgoing edges. For unweighted graphs, all edge
         weights are considered to be 1.0.
 
+        Parameters
+        ----------
+        of : Relationship, optional
+            A unary relationship containing a subset of the graph's nodes. When
+            provided, constrains the domain of the weighted degree computation: only
+            weighted degrees of nodes in this relationship are computed and returned.
+
         Returns
         -------
         Relationship
@@ -2508,7 +2515,7 @@ class Graph():
 
         Examples
         --------
-        >>> from relationalai.semantics import Model, define, select, Float
+        >>> from relationalai.semantics import Model, define, select, where, union, Float
         >>> from relationalai.semantics.reasoners.graph import Graph
         >>>
         >>> # 1. Set up a directed, weighted graph
@@ -2528,12 +2535,42 @@ class Graph():
         >>> # 3. Select the weighted degree of each node and inspect
         >>> node, node_weighted_degree = Node.ref("node"), Float.ref("node_weighted_degree")
         >>> weighted_degree = graph.weighted_degree()
-        >>> select(node.id, node_weighted_degree).where(weighted_degree(node, node_weighted_degree)).inspect()
+        >>> select(
+        ...     node.id, node_weighted_degree
+        ... ).where(
+        ...     weighted_degree(node, node_weighted_degree)
+        ... ).inspect()
         ▰▰▰▰ Setup complete
             id  node_weighted_degree
         0   1                   0.0
         1   2                   1.0
         2   3                   1.0
+        >>>
+        >>> # 4. Use 'of' parameter to constrain the set of nodes to compute weighted degree of
+        >>> subset = model.Relationship(f"{{node:{Node}}} is in subset")
+        >>> node = Node.ref()
+        >>> where(union(node.id == 2, node.id == 3)).define(subset(node))
+        >>> constrained_weighted_degree = graph.weighted_degree(of=subset)
+        >>> select(
+        ...     node.id, node_weighted_degree
+        ... ).where(
+        ...     constrained_weighted_degree(node, node_weighted_degree)
+        ... ).inspect()
+        ▰▰▰▰ Setup complete
+            id  node_weighted_degree
+        0   2                   1.0
+        1   3                   1.0
+
+        Notes
+        -----
+        The ``weighted_degree()`` method, called with no parameters, computes and caches
+        the full weighted degree relationship, providing efficient reuse across multiple
+        calls to ``weighted_degree()``. In contrast, ``weighted_degree(of=subset)`` computes a
+        constrained relationship specific to the passed-in ``subset`` and that
+        call site. When a significant fraction of the weighted degree relation is needed
+        across a program, ``weighted_degree()`` is typically more efficient; this is the
+        typical case. Use ``weighted_degree(of=subset)`` only when small subsets of the
+        weighted degree relationship are needed collectively across the program.
 
         See Also
         --------
@@ -2541,23 +2578,55 @@ class Graph():
         weighted_outdegree
 
         """
-        return self._weighted_degree
+        if of is None:
+            return self._weighted_degree
+        else:
+            # Validate the 'of' parameter
+            self._validate_node_subset_parameter(of)
+            return self._weighted_degree_of(of)
 
     @cached_property
     def _weighted_degree(self):
         """Lazily define and cache the self._weighted_degree relationship."""
+        return self._create_weighted_degree_relationship(nodes_subset=None)
+
+    def _weighted_degree_of(self, nodes_subset: Relationship):
+        """
+        Create a weighted degree relationship constrained to the subset of nodes
+        in `nodes_subset`. Note this relationship is not cached; it is
+        specific to the callsite.
+        """
+        return self._create_weighted_degree_relationship(nodes_subset=nodes_subset)
+
+    def _create_weighted_degree_relationship(self, *, nodes_subset: Optional[Relationship]):
         _weighted_degree_rel = self._model.Relationship(f"{{node:{self._NodeConceptStr}}} has weighted degree {{weight:Float}}")
 
         if self.directed:
+            # For directed graphs, weighted degree is the sum of weighted indegree and weighted outdegree.
+            if nodes_subset is None:
+                weighted_indegree_rel = self._weighted_indegree
+                weighted_outdegree_rel = self._weighted_outdegree
+            else:
+                weighted_indegree_rel = self._weighted_indegree_of(nodes_subset)
+                weighted_outdegree_rel = self._weighted_outdegree_of(nodes_subset)
+
             inweight, outweight = Float.ref(), Float.ref()
             where(
-                self._weighted_indegree(self.Node, inweight),
-                self._weighted_outdegree(self.Node, outweight),
+                weighted_indegree_rel(self.Node, inweight),
+                weighted_outdegree_rel(self.Node, outweight),
             ).define(_weighted_degree_rel(self.Node, inweight + outweight))
         elif not self.directed:
+            # Choose the appropriate node set
+            if nodes_subset is None:
+                # No constraint - use all nodes
+                node_set = self.Node
+            else:
+                # Constrained to nodes in the subset
+                node_set = nodes_subset
+
             dst, weight = self.Node.ref(), Float.ref()
             where(
-                self.Node,
+                node_set(self.Node),
                 _weighted_degree := sum(dst, weight).per(self.Node).where(self._weight(self.Node, dst, weight)) | 0.0,
             ).define(_weighted_degree_rel(self.Node, _weighted_degree))
 
@@ -2565,13 +2634,20 @@ class Graph():
 
 
     @include_in_docs
-    def weighted_indegree(self):
+    def weighted_indegree(self, *, of: Optional[Relationship] = None):
         """Returns a binary relationship containing the weighted indegree of each node.
 
         A node's weighted indegree is the sum of the weights of all incoming
         edges. For undirected graphs, this is identical to `weighted_degree`.
         For unweighted graphs, all edge weights are considered to be 1.0.
 
+        Parameters
+        ----------
+        of : Relationship, optional
+            A unary relationship containing a subset of the graph's nodes. When
+            provided, constrains the domain of the weighted indegree computation: only
+            weighted indegrees of nodes in this relationship are computed and returned.
+
         Returns
         -------
         Relationship
@@ -2596,7 +2672,7 @@ class Graph():
 
         Examples
         --------
-        >>> from relationalai.semantics import Model, define, select, Float
+        >>> from relationalai.semantics import Model, define, select, where, union, Float
         >>> from relationalai.semantics.reasoners.graph import Graph
         >>>
         >>> # 1. Set up a directed, weighted graph
@@ -2626,6 +2702,28 @@ class Graph():
         0   1                    -1.0
         1   2                     1.0
         2   3                     1.0
+        >>>
+        >>> # 4. Use 'of' parameter to constrain the set of nodes to compute weighted indegree of
+        >>> subset = model.Relationship(f"{{node:{Node}}} is in subset")
+        >>> node = Node.ref()
+        >>> where(union(node.id == 2, node.id == 3)).define(subset(node))
+        >>> constrained_weighted_indegree = graph.weighted_indegree(of=subset)
+        >>> select(node.id, node_weighted_indegree).where(constrained_weighted_indegree(node, node_weighted_indegree)).inspect()
+        ▰▰▰▰ Setup complete
+            id  node_weighted_indegree
+        0   2                     1.0
+        1   3                     1.0
+
+        Notes
+        -----
+        The ``weighted_indegree()`` method, called with no parameters, computes and caches
+        the full weighted indegree relationship, providing efficient reuse across multiple
+        calls to ``weighted_indegree()``. In contrast, ``weighted_indegree(of=subset)`` computes a
+        constrained relationship specific to the passed-in ``subset`` and that
+        call site. When a significant fraction of the weighted indegree relation is needed
+        across a program, ``weighted_indegree()`` is typically more efficient; this is the
+        typical case. Use ``weighted_indegree(of=subset)`` only when small subsets of the
+        weighted indegree relationship are needed collectively across the program.
 
         See Also
         --------
@@ -2633,16 +2731,49 @@ class Graph():
         weighted_outdegree
 
         """
-        return self._weighted_indegree
+        # TODO: It looks like the weights in the example in the docstring above
+        #   are holdovers from a version of the library that did not disallow
+        #   negative weights. Need to update the example to use only non-negative weights.
+        if of is None:
+            return self._weighted_indegree
+        else:
+            # Validate the 'of' parameter
+            self._validate_node_subset_parameter(of)
+            return self._weighted_indegree_of(of)
 
     @cached_property
     def _weighted_indegree(self):
         """Lazily define and cache the self._weighted_indegree relationship."""
+        return self._create_weighted_indegree_relationship(nodes_subset=None)
+
+    def _weighted_indegree_of(self, nodes_subset: Relationship):
+        """
+        Create a weighted indegree relationship constrained to the subset of nodes
+        in `nodes_subset`. Note this relationship is not cached; it is
+        specific to the callsite.
+        """
+        return self._create_weighted_indegree_relationship(nodes_subset=nodes_subset)
+
+    def _create_weighted_indegree_relationship(self, *, nodes_subset: Optional[Relationship]):
         _weighted_indegree_rel = self._model.Relationship(f"{{node:{self._NodeConceptStr}}} has weighted indegree {{weight:Float}}")
 
+        # Choose the appropriate node set
+        if nodes_subset is None:
+            # No constraint - use all nodes
+            node_set = self.Node
+        else:
+            # Constrained to nodes in the subset
+            node_set = nodes_subset
+        # TODO: In a future cleanup pass, replace `node_set` with a `node_constraint`
+        #   that replaces the `node_set(self.Node)` in the where clause below,
+        #   and generates only `self.Node` (rather than `self.Node(self.Node)`)
+        #   in the `subset is None` case. This applies to a couple other
+        #   degree-of type relations as well.
+
+        # Apply the weighted indegree logic for both cases
         src, inweight = self.Node.ref(), Float.ref()
         where(
-            self.Node,
+            node_set(self.Node),
             _weighted_indegree := sum(src, inweight).per(self.Node).where(self._weight(src, self.Node, inweight)) | 0.0,
         ).define(_weighted_indegree_rel(self.Node, _weighted_indegree))
 
@@ -2650,13 +2781,20 @@ class Graph():
 
 
     @include_in_docs
-    def weighted_outdegree(self):
+    def weighted_outdegree(self, *, of: Optional[Relationship] = None):
         """Returns a binary relationship containing the weighted outdegree of each node.
 
         A node's weighted outdegree is the sum of the weights of all outgoing
         edges. For undirected graphs, this is identical to `weighted_degree`.
         For unweighted graphs, all edge weights are considered to be 1.0.
 
+        Parameters
+        ----------
+        of : Relationship, optional
+            A unary relationship containing a subset of the graph's nodes. When
+            provided, constrains the domain of the weighted outdegree computation: only
+            weighted outdegrees of nodes in this relationship are computed and returned.
+
         Returns
         -------
         Relationship
@@ -2681,7 +2819,7 @@ class Graph():
 
         Examples
         --------
-        >>> from relationalai.semantics import Model, define, select, Float
+        >>> from relationalai.semantics import Model, define, select, where, union, Float
         >>> from relationalai.semantics.reasoners.graph import Graph
         >>>
         >>> # 1. Set up a directed, weighted graph
@@ -2711,6 +2849,32 @@ class Graph():
         0   1                      1.0
         1   2                      0.0
         2   3                      0.0
+        >>>
+        >>> # 4. Use 'of' parameter to constrain the set of nodes to compute weighted outdegree of
+        >>> subset = model.Relationship(f"{{node:{Node}}} is in subset")
+        >>> node = Node.ref()
+        >>> where(union(node.id == 1, node.id == 2)).define(subset(node))
+        >>> constrained_weighted_outdegree = graph.weighted_outdegree(of=subset)
+        >>> select(
+        ...     node.id, node_weighted_outdegree
+        ... ).where(
+        ...     constrained_weighted_outdegree(node, node_weighted_outdegree)
+        ... ).inspect()
+        ▰▰▰▰ Setup complete
+            id  node_weighted_outdegree
+        0   1                      1.0
+        1   2                      0.0
+
+        Notes
+        -----
+        The ``weighted_outdegree()`` method, called with no parameters, computes and caches
+        the full weighted outdegree relationship, providing efficient reuse across multiple
+        calls to ``weighted_outdegree()``. In contrast, ``weighted_outdegree(of=subset)`` computes a
+        constrained relationship specific to the passed-in ``subset`` and that
+        call site. When a significant fraction of the weighted outdegree relation is needed
+        across a program, ``weighted_outdegree()`` is typically more efficient; this is the
+        typical case. Use ``weighted_outdegree(of=subset)`` only when small subsets of the
+        weighted outdegree relationship are needed collectively across the program.
 
         See Also
         --------
@@ -2718,16 +2882,41 @@ class Graph():
         weighted_indegree
 
         """
-        return self._weighted_outdegree
+        if of is None:
+            return self._weighted_outdegree
+        else:
+            # Validate the 'of' parameter
+            self._validate_node_subset_parameter(of)
+            return self._weighted_outdegree_of(of)
 
     @cached_property
     def _weighted_outdegree(self):
         """Lazily define and cache the self._weighted_outdegree relationship."""
+        return self._create_weighted_outdegree_relationship(nodes_subset=None)
+
+    def _weighted_outdegree_of(self, nodes_subset: Relationship):
+        """
+        Create a weighted outdegree relationship constrained to the subset of nodes
+        in `nodes_subset`. Note this relationship is not cached; it is
+        specific to the callsite.
+        """
+        return self._create_weighted_outdegree_relationship(nodes_subset=nodes_subset)
+
+    def _create_weighted_outdegree_relationship(self, *, nodes_subset: Optional[Relationship]):
         _weighted_outdegree_rel = self._model.Relationship(f"{{node:{self._NodeConceptStr}}} has weighted outdegree {{weight:Float}}")
 
+        # Choose the appropriate node set
+        if nodes_subset is None:
+            # No constraint - use all nodes
+            node_set = self.Node
+        else:
+            # Constrained to nodes in the subset
+            node_set = nodes_subset
+
+        # Apply the weighted outdegree logic for both cases
         dst, outweight = self.Node.ref(), Float.ref()
         where(
-            self.Node,
+            node_set(self.Node),
             _weighted_outdegree := sum(dst, outweight).per(self.Node).where(self._weight(self.Node, dst, outweight)) | 0.0,
         ).define(_weighted_outdegree_rel(self.Node, _weighted_outdegree))
 
@@ -2735,14 +2924,24 @@ class Graph():
 
 
     @include_in_docs
-    def degree_centrality(self):
+    def degree_centrality(self, *, of: Optional[Relationship] = None):
         """Returns a binary relationship containing the degree centrality of each node.
 
         Degree centrality is a measure of a node's importance, defined as its
         degree (or weighted degree for weighted graphs) divided by the number
-        of other nodes in the graph. For simple graphs without self-loops, this
-        value will be at most 1.0; graphs with self-loops might have nodes
-        with a degree centrality greater than 1.0.
+        of other nodes in the graph.
+
+        For unewighted graphs without self-loops, this value will be at most 1.0;
+        unweighted graphs with self-loops might have nodes with a degree centrality
+        greater than 1.0. Weighted graphs may have degree centralities
+        greater than 1.0 as well.
+
+        Parameters
+        ----------
+        of : Relationship, optional
+            A unary relationship containing a subset of the graph's nodes. When
+            provided, constrains the domain of the degree centrality computation: only
+            degree centralities of nodes in this relationship are computed and returned.
 
         Returns
         -------
@@ -2798,6 +2997,20 @@ class Graph():
         2   3   1.000000
         3   4   0.666667
 
+        >>> # 4. Use 'of' parameter to constrain the set of nodes to compute degree centrality of
+        >>> # Define a subset containing only nodes 2 and 3
+        >>> subset = model.Relationship(f"{{node:{Node}}} is in subset")
+        >>> node = Node.ref()
+        >>> where(union(node.id == 2, node.id == 3)).define(subset(node))
+        >>>
+        >>> # Get degree centralities only of nodes in the subset
+        >>> constrained_degree_centrality = graph.degree_centrality(of=subset)
+        >>> select(node.id, centrality).where(constrained_degree_centrality(node, centrality)).inspect()
+        ▰▰▰▰ Setup complete
+           id centrality
+        0   2        1.0
+        1   3        1.0
+
         **Weighted Graph Example**
 
         >>> from relationalai.semantics import Model, define, select, Float
@@ -2827,52 +3040,90 @@ class Graph():
         1   2        1.75
         2   3        1.00
 
+        Notes
+        -----
+        The ``degree_centrality()`` method, called with no parameters, computes and caches
+        the full degree centrality relationship, providing efficient reuse across multiple
+        calls to ``degree_centrality()``. In contrast, ``degree_centrality(of=subset)`` computes a
+        constrained relationship specific to the passed-in ``subset`` and that
+        call site. When a significant fraction of the degree centrality relation is needed
+        across a program, ``degree_centrality()`` is typically more efficient; this is the
+        typical case. Use ``degree_centrality(of=subset)`` only when small subsets of the
+        degree centrality relationship are needed collectively across the program.
+
         See Also
         --------
         degree
         weighted_degree
 
         """
-        return self._degree_centrality
+        if of is None:
+            return self._degree_centrality
+        else:
+            # Validate the 'of' parameter
+            self._validate_node_subset_parameter(of)
+            return self._degree_centrality_of(of)
 
     @cached_property
     def _degree_centrality(self):
         """Lazily define and cache the self._degree_centrality relationship."""
+        return self._create_degree_centrality_relationship(nodes_subset=None)
+
+    def _degree_centrality_of(self, nodes_subset: Relationship):
+        """
+        Create a degree centrality relationship constrained to the subset of nodes
+        in `nodes_subset`. Note this relationship is not cached; it is
+        specific to the callsite.
+        """
+        return self._create_degree_centrality_relationship(nodes_subset=nodes_subset)
+
+    def _create_degree_centrality_relationship(self, *, nodes_subset: Optional[Relationship]):
+        """Create a degree centrality relationship, optionally constrained to a subset of nodes."""
         _degree_centrality_rel = self._model.Relationship(f"{{node:{self._NodeConceptStr}}} has {{degree_centrality:Float}}")
 
-        degree = Integer.ref()
-        weighted_degree = Float.ref()
+        if nodes_subset is None:
+            degree_rel = self._degree
+            node_constraint = [] # No constraint on nodes.
+        else:
+            degree_rel = self._degree_of(nodes_subset)
+            node_constraint = [nodes_subset(self.Node)]  # Nodes constrained to given subset.
 
-        # TODO: Collapse the logic for these edge cases using a match.
+        degree = Integer.ref()
 
         # A single isolated node has degree centrality zero.
         where(
             self._num_nodes(1),
-            self._degree(self.Node, 0)
+            *node_constraint,
+            degree_rel(self.Node, 0)
         ).define(_degree_centrality_rel(self.Node, 0.0))
 
         # A single non-isolated node has degree centrality one.
         where(
             self._num_nodes(1),
-            self._degree(self.Node, degree),
+            *node_constraint,
+            degree_rel(self.Node, degree),
             degree > 0
         ).define(_degree_centrality_rel(self.Node, 1.0))
 
         # General case, i.e. with more than one node.
-        num_nodes = Integer.ref()
         if self.weighted:
-            where(
-                self._num_nodes(num_nodes),
-                num_nodes > 1,
-                self._weighted_degree(self.Node, weighted_degree)
-            ).define(_degree_centrality_rel(self.Node, weighted_degree / (num_nodes - 1.0)))
-        elif not self.weighted:
-            where(
-                self._num_nodes(num_nodes),
-                num_nodes > 1,
-                self._degree(self.Node, degree)
-            ).define(_degree_centrality_rel(self.Node, degree / (num_nodes - 1.0)))
-            # TODO: .to_df() second-column dtype becomes `object` rather than `Float`.
+            maybe_weighted_degree = Float.ref()
+            if nodes_subset is None:
+                maybe_weighted_degree_rel = self._weighted_degree
+            else:
+                maybe_weighted_degree_rel = self._weighted_degree_of(nodes_subset)
+        else: # not self.weighted
+            maybe_weighted_degree = Integer.ref()
+            maybe_weighted_degree_rel = degree_rel
+
+        num_nodes = Integer.ref()
+
+        where(
+            self._num_nodes(num_nodes),
+            num_nodes > 1,
+            *node_constraint,
+            maybe_weighted_degree_rel(self.Node, maybe_weighted_degree)
+        ).define(_degree_centrality_rel(self.Node, maybe_weighted_degree / (num_nodes - 1.0)))
 
         return _degree_centrality_rel
 
@@ -3614,12 +3865,19 @@ class Graph():
 
 
     @include_in_docs
-    def triangle_count(self):
+    def triangle_count(self, *, of: Optional[Relationship] = None):
         """Returns a binary relationship containing the number of unique triangles each node belongs to.
 
         A triangle is a set of three nodes where each node has a directed
         or undirected edge to the other two nodes, forming a 3-cycle.
 
+        Parameters
+        ----------
+        of : Relationship, optional
+            A unary relationship containing a subset of the graph's nodes. When
+            provided, constrains the domain of the triangle count computation: only
+            triangle counts of nodes in this relationship are computed and returned.
+
         Returns
         -------
         Relationship
@@ -3675,6 +3933,31 @@ class Graph():
         3   4      0
         4   5      0
 
+        >>> # 4. Use 'of' parameter to constrain the set of nodes to compute triangle counts of
+        >>> # Define a subset containing only nodes 1 and 3
+        >>> subset = model.Relationship(f"{{node:{Node}}} is in subset")
+        >>> node = Node.ref()
+        >>> where(union(node.id == 1, node.id == 3)).define(subset(node))
+        >>>
+        >>> # Get triangle counts only of nodes in the subset
+        >>> constrained_triangle_count = graph.triangle_count(of=subset)
+        >>> select(node.id, count).where(constrained_triangle_count(node, count)).inspect()
+        ▰▰▰▰ Setup complete
+           id  count
+        0   1      1
+        1   3      1
+
+        Notes
+        -----
+        The ``triangle_count()`` method, called with no parameters, computes and caches
+        the full triangle count relationship, providing efficient reuse across multiple
+        calls to ``triangle_count()``. In contrast, ``triangle_count(of=subset)`` computes a
+        constrained relationship specific to the passed-in ``subset`` and that
+        call site. When a significant fraction of the triangle count relation is needed
+        across a program, ``triangle_count()`` is typically more efficient; this is the
+        typical case. Use ``triangle_count(of=subset)`` only when small subsets of the
+        triangle count relationship are needed collectively across the program.
+
         See Also
         --------
         triangle
@@ -3682,15 +3965,35 @@ class Graph():
         num_triangles
 
         """
+        if of is not None:
+            self._validate_node_subset_parameter(of)
+            return self._triangle_count_of(of)
         return self._triangle_count
 
     @cached_property
     def _triangle_count(self):
         """Lazily define and cache the self._triangle_count relationship."""
+        return self._create_triangle_count_relationship(nodes_subset=None)
+
+    def _triangle_count_of(self, nodes_subset: Relationship):
+        """
+        Create a triangle count relationship constrained to the subset of nodes
+        in `nodes_subset`. Note this relationship is not cached; it is
+        specific to the callsite.
+        """
+        return self._create_triangle_count_relationship(nodes_subset=nodes_subset)
+
+    def _create_triangle_count_relationship(self, *, nodes_subset: Optional[Relationship]):
+        """Create a triangle count relationship, optionally constrained to a subset of nodes."""
         _triangle_count_rel = self._model.Relationship(f"{{node:{self._NodeConceptStr}}} belongs to {{count:Integer}} triangles")
 
+        if nodes_subset is None:
+            node_constraint = self.Node # No constraint on nodes.
+        else:
+            node_constraint = nodes_subset(self.Node)  # Nodes constrained to given subset.
+
         where(
-            self.Node,
+            node_constraint,
             _count := self._nonzero_triangle_count_fragment(self.Node) | 0
         ).define(_triangle_count_rel(self.Node, _count))
 

relationalai/semantics/reasoners/optimization/solvers_dev.py

@@ -2,12 +2,15 @@ from __future__ import annotations
 from typing import Union
 import textwrap
 import uuid
+import time
 
 from relationalai.semantics.snowflake import Table
 from relationalai.semantics import std
 from relationalai.semantics.internal import internal as b # TODO(coey) change b name or remove b.?
 from relationalai.semantics.rel.executor import RelExecutor
 from relationalai.semantics.lqp.executor import LQPExecutor
+from relationalai.tools.constants import DEFAULT_QUERY_TIMEOUT_MINS
+from relationalai.util.timeout import calc_remaining_timeout_minutes
 
 from .common import make_name
 from relationalai.experimental.solvers import Solver
@@ -243,6 +246,17 @@ class SolverModelDev:
         app_name = resources.get_app_name()
         print(app_name)
 
+        # Note: currently the query timeout is not propagated to the steps 'export model
+        # relations', and 'import result relations'. For those steps the default query
+        # timeout value defined in the config will apply.
+        # TODO: propagate the query timeout to those steps as well.
+        query_timeout_mins = kwargs.get("query_timeout_mins", None)
+        config = self._model._config
+        if query_timeout_mins is None and (timeout_value := config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS)) is not None:
+            query_timeout_mins = int(timeout_value)
+        config_file_path = getattr(config, 'file_path', None)
+        start_time = time.monotonic()
+
         # 1. export model relations
         print("export model relations")
         # TODO(coey) perf: only export the relations that are actually used in the model
@@ -266,6 +280,9 @@ class SolverModelDev:
             b.select(*rel._field_refs).where(rel(*rel._field_refs)).into(table)
 
         # 2. execute solver job and wait for completion
+        remaining_timeout_minutes = calc_remaining_timeout_minutes(
+            start_time, query_timeout_mins, config_file_path=config_file_path,
+        )
         print("execute solver job")
         payload = {
             "solver": solver.solver_name.lower(),
@@ -273,7 +290,9 @@ class SolverModelDev:
             "input_id": input_id,
             "data_type": self._data_type
         }
-        job_id = solver._exec_job(payload, log_to_console=log_to_console)
+        job_id = solver._exec_job(
+            payload, log_to_console=log_to_console, query_timeout_mins=remaining_timeout_minutes,
+        )
         print(f"job id: {job_id}") # TODO(coey) maybe job_id is not useful
 
         # 3. import result relations