relationalai 0.13.5__py3-none-any.whl → 1.0.0a1__py3-none-any.whl

relationalai/experimental/solvers.py

@@ -1,1095 +0,0 @@
-from __future__ import annotations
-import time
-from typing import Any, List, Optional, cast
-from dataclasses import dataclass
-import textwrap
-from .. import dsl, std
-from ..std import rel
-from ..metamodel import Builtins
-from ..tools.cli_controls import Spinner
-from ..tools.constants import DEFAULT_QUERY_TIMEOUT_MINS
-from .. import debugging
-from .. errors import ResponseStatusException
-import uuid
-import relationalai
-import json
-from ..clients.util import poll_with_specified_overhead
-from ..clients.resources.snowflake import Resources as SnowflakeResources, APP_NAME
-from ..clients.resources.snowflake import DirectAccessResources
-from ..clients.direct_access_client import DirectAccessClient
-from ..util.timeout import calc_remaining_timeout_minutes
-
-rel_sv = rel._tagged(Builtins.SingleValued)
-
-ENGINE_TYPE_SOLVER = "SOLVER"
-# TODO (dba) The ERP still uses `worker` instead of `engine`. Change
-# this once we fix this in the ERP.
-WORKER_ERRORS = ["worker is suspended", "create/resume", "worker not found", "no workers found", "worker was deleted"]
-ENGINE_ERRORS = ["engine is suspended", "create/resume", "engine not found", "no engines found", "engine was deleted"]
-ENGINE_NOT_READY_MSGS = ["worker is in pending", "worker is provisioning", "worker is not ready to accept jobs"]
-
-# --------------------------------------------------
-# SolverModel object.
-# --------------------------------------------------
-
-class SolverModel:
-    def __init__(self, graph:dsl.Graph):
-        self.graph = graph
-        self.id = dsl.next_id()
-        self.scope = f"solvermodel{self.id}_"
-        scope = self.scope
-        self.Variable = dsl.Type(graph, "variables", scope=scope)
-        self.MinObjective = dsl.Type(graph, "min_objectives", scope=scope)
-        self.MaxObjective = dsl.Type(graph, "max_objectives", scope=scope)
-        self.Constraint = dsl.Type(graph, "constraints", scope=scope)
-        self.Solution = dsl.Type(graph, "solutions", scope=scope)
-        self.components = [
-            (self.MinObjective, "minimization objectives"),
-            (self.MaxObjective, "maximization objectives"),
-            (self.Constraint, "constraints"),
-        ]
-        self.solve_index = 0
-        self.is_solved = False
-
-        # Install model helpers.
-        self.graph.install_raw(textwrap.dedent(f"""
-            @inline
-            def _solverlib_ho_appl(op, {{R}}, s): rel_primitive_solverlib_ho_appl(R, op, s)
-
-            @inline
-            def _solver_unwrap({{R}}, h, x...): exists((v) | R(v, x...) and pyrel_unwrap(v, h))
-
-            declare {scope}variable_name
-            declare {scope}component_name
-            declare {scope}serialized
-            declare {scope}primal_start
-
-            def {scope}component_string(h, s):
-                rel_primitive_solverlib_print_expr(
-                    {scope}serialized[h], _solver_unwrap[{scope}variable_name], s
-                )
-
-            declare {scope}solve_output
-        """))
-        return None
-
-    # Add an entity to the variable set, (optionally) set a string name from the
-    # arguments and add domain constraints on the variable.
-    def variable(
-        self,
-        var, # variable entity
-        name_args:List|None=None, # list of strings to concatenate into a string name
-        type:str|None=None, # variable type: "integer" or "zero_one"
-        lower:int|float|None=None, # lower bound
-        upper:int|float|None=None, # upper bound
-        fixed:int|float|None=None, # fixed value
-        start:int|float|None=None, # (primal) start value
-    ):
-        if type not in {"integer", "zero_one", None}:
-            raise Exception(f"Invalid domain type: {type}.")
-        var.set(self.Variable)
-
-        # Set variable name.
-        if name_args:
-            var.set(**{f"{self.scope}variable_name": make_string(name_args)})
-
-        # Add domain constraints.
-        cons = []
-        if fixed is not None:
-            cons.append(eq(var, fixed))
-        if type == "zero_one":
-            cons.append(zero_one(var))
-        if lower is not None and upper is not None:
-            if type == "integer":
-                cons.append(integer_interval(var, lower, upper))
-            else:
-                cons.append(interval(var, lower, upper))
-        else:
-            if type == "integer":
-                cons.append(integer(var))
-            if lower is not None:
-                cons.append(gte(var, lower))
-            if upper is not None:
-                cons.append(lte(var, upper))
-        if len(cons) == 1:
-            self.constraint(cons[0])
-        elif len(cons) > 1:
-            self.constraint(and_(*cons))
-
-        # Set primal start.
-        if start is not None:
-            var.set(**{f"{self.scope}primal_start": start})
-        return var
-
-    # Get variable string name.
-    def variable_name(self, var):
-        return std.alias(getattr(var, f"{self.scope}variable_name"), "name")
-
-    # Add a constraint, minimization objective, or maximization objective.
-    def constraint(self, expr, name_args:List|None=None):
-        return self._add_component(self.Constraint, expr, name_args)
-
-    def min_objective(self, expr, name_args:List|None=None):
-        return self._add_component(self.MinObjective, expr, name_args)
-
-    def max_objective(self, expr, name_args:List|None=None):
-        return self._add_component(self.MaxObjective, expr, name_args)
-
-    def _add_component(self, typ, expr, name_args:List|None):
-        comp = typ.add(serialized=_wrap_expr(expr))
-        if name_args:
-            comp.set(component_name=make_string(name_args))
-        return comp
-
-    # Get component string name.
-    def component_name(self, comp):
-        return std.alias(comp.component_name, "name")
-
-    # Get serialized component string in human-readable format.
-    def component_string(self, comp):
-        return std.alias(comp.component_string, "string")
-
-    # Summarize the model by printing the number of variables and components.
-    # Use outside a rule/query.
-    def summarize(self):
-        with self.graph.query() as select:
-            vars = select(std.aggregates.count(self.Variable()))
-        s = f"Model has: {vars.results.iat[0, 0]} variables"
-        for (c_type, c_name) in self.components:
-            with self.graph.query() as select:
-                exprs = select(std.aggregates.count(c_type()))
-            if not exprs.results.empty:
-                s += f", {exprs.results.iat[0, 0]} {c_name}"
-        print(s)
-        return None
-
-    # Print the model in human-readable format. Use outside a rule/query.
-    def print(self):
-        with self.graph.query() as select:
-            vars = select(rel.last(getattr(rel, f"{self.scope}variable_name")))
-        print("variables:")
-        print(vars.results.to_string(index=False, header=False))
-        for (c_type, c_name) in self.components:
-            with self.graph.query() as select:
-                exprs = select(self.component_string(c_type()))
-            if not exprs.results.empty:
-                print(c_name + ":")
-                print(exprs.results.to_string(index=False, header=False))
-        return None
-
-    # Solve the model given a solver and solver options. Use outside a rule/query.
-    def solve(self, solver: Solver, log_to_console=True, **kwargs):
-        self.is_solved = False
-        self.solve_index += 1
-        self.solve_output = dsl.RelationNS([f"{self.scope}solve_output"], f"i_{self.solve_index}")
-
-        options = kwargs
-        options["version"] = 1
-
-        # Validate options.
-        for k, v in options.items():
-            if not isinstance(k, str):
-                raise Exception(f"Invalid parameter key. Expected string, got {type(k)} for {k}.")
-            if not isinstance(v, (int, float, str, bool)):
-                raise Exception(
-                    f"Invalid parameter value. Expected string, integer, float, or boolean, got {type(v)} for {k}."
-                )
-
-        # Run the solve query and insert the solve_output result.
-        scope = self.scope
-        variable_name_string = f"{scope}variable_name" if "print_format" in options else "{}"
-        component_name_string = f"{scope}component_name" if "print_format" in options else "{}"
-
-        input_id = uuid.uuid4()
-        model_uri = f"snowflake://APP_STATE.RAI_INTERNAL_STAGE/job-inputs/solver/{input_id}/model.binpb"
-        sf_input_uri = f"snowflake://job-inputs/solver/{input_id}/model.binpb"
-
-        payload: dict[str, Any] = {"solver": solver.solver_name.lower()}
-        payload["options"] = options
-        payload["model_uri"] = sf_input_uri
-
-        rai_config = self.graph._config
-        query_timeout_mins = kwargs.get("query_timeout_mins", None)
-        if query_timeout_mins is None and (timeout_value := rai_config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS)) is not None:
-            query_timeout_mins = int(timeout_value)
-        config_file_path = getattr(rai_config, 'file_path', None)
-        start_time = time.monotonic()
-        remaining_timeout_minutes = query_timeout_mins
-        response = self.graph.exec_raw(
-            textwrap.dedent(f"""
-            @inline
-            def {scope}specialized_components(t, h, s):
-                exists((v) | {{
-                    (:min_objective, {scope}min_objectives);
-                    (:max_objective, {scope}max_objectives);
-                    (:constraint, {scope}constraints);
-                    }}(t, v) and pyrel_unwrap(v, h) and {scope}serialized(v, s)
-                )
-
-            @no_diagnostics(:EXPERIMENTAL)
-            def {scope}model_string {{
-                rel_primitive_solverlib_model_string[{{
-                    (:variable, _solver_unwrap[{scope}variables]);
-                    {scope}specialized_components;
-                    (:variable_name, _solver_unwrap[{variable_name_string}]);
-                    (:expression_name, _solver_unwrap[{component_name_string}]);
-                    (:primal_start, _solver_unwrap[{scope}primal_start]);
-                }}]
-            }}
-
-            ic model_not_empty("Solver model is empty.") requires not empty({scope}model_string)
-
-            def config[:envelope, :content_type]: "application/octet-stream"
-            def config[:envelope, :payload, :data]: {scope}model_string
-            def config[:envelope, :payload, :path]: "{model_uri}"
-            def export {{ config }}
-            """),
-            query_timeout_mins=remaining_timeout_minutes,
-        )
-        txn = response.transaction or {}
-        # The above `exec_raw` will throw an error if the transaction
-        # gets aborted. But in the case it gets cancelled, by the user
-        # or the system, it won't throw. In that case we also did not
-        # upload the input model.
-        if txn["state"] == "":
-            txn = solver.provider.resources.get_transaction(txn["id"]) or {}
-        if txn["state"] != "COMPLETED":
-            raise Exception(f"Transaction that materializes the solver inputs did not complete! ID: `{txn['id']}` State `{txn['state']}`")
-
-        # 2. Execute job and wait for completion.
-        remaining_timeout_minutes = calc_remaining_timeout_minutes(
-            start_time, query_timeout_mins, config_file_path=config_file_path
-        )
-        job_id = solver._exec_job(payload, log_to_console=log_to_console, query_timeout_mins=remaining_timeout_minutes)
-
-        # 3. Extract result.
-        remaining_timeout_minutes = calc_remaining_timeout_minutes(
-            start_time, query_timeout_mins, config_file_path=config_file_path
-        )
-        res = self.graph.exec_raw(
-            textwrap.dedent(f"""
-            ic result_not_empty("Solver result is empty.") requires not empty(result)
-
-            def result {{load_binary["snowflake://APP_STATE.RAI_INTERNAL_STAGE/job-results/{job_id}/result.binpb"] }}
-            def delete[:{scope}solve_output, :"i_{self.solve_index}"]: {scope}solve_output[:"i_{self.solve_index}"]
-
-            @no_diagnostics(:EXPERIMENTAL)
-            def insert[:{scope}solve_output, :"i_{self.solve_index}"]:
-                rel_primitive_solverlib_extract[result]
-
-            def output[:solver_error]: {scope}solve_output[:"i_{self.solve_index}", :error]
-            """),
-            readonly=False,
-            query_timeout_mins=remaining_timeout_minutes,
-        )
-        errors = []
-        for result in res.results:
-            if result["relationId"] == "/:output/:solver_error/String":
-                errors.extend(result["table"]["v1"])
-
-        # 4. Map results to solution.
-        with self.graph.rule(dynamic=True):
-            sol = self.Solution.add(index = self.solve_index)
-            for name in {"error", "termination_status", "solve_time_sec", "objective_value", "solver_version", "printed_model"}:
-                val = dsl.create_var()
-                getattr(self.solve_output, name)(val)
-                sol.set(**{name:val})
-
-        self.is_solved = True
-
-        if len(errors) > 0:
-            raise Exception("\n".join(errors))
-
-        return None
-
-    # Get scalar result information after solving.
-    def __getattr__(self, name:str):
-        if not self.is_solved:
-            raise Exception("Model has not been solved yet.")
-        if name in {"error", "termination_status", "solve_time_sec", "objective_value", "solver_version", "printed_model"}:
-            return getattr(self.Solution(index = self.solve_index), name)
-        else:
-            return None
-
-    # Get variable point values after solving. If `index` is specified, get the value
-    # of the variable in the return the `index`-th solution.
-    def value(self, var, index:int|None=None):
-        if not self.is_solved:
-            raise Exception("Model has not been solved yet.")
-        val = dsl.create_var()
-        unwrap_var = rel_sv.pyrel_unwrap(var)
-        if index:
-            self.solve_output.points(unwrap_var, index, val)
-        else:
-            self.solve_output.point(unwrap_var, val)
-        std.alias(val, "value")
-        return val
-
-# --------------------------------------------------
-# Operator definitions
-# --------------------------------------------------
-
-# Builtin binary operators
-
-def plus(left, right):
-    return _make_fo_expr(10, left, right)
-
-def minus(left, right):
-    return _make_fo_expr(11, left, right)
-
-def mult(left, right):
-    return _make_fo_expr(12, left, right)
-
-def div(left, right):
-    return _make_fo_expr(13, left, right)
-
-def pow(left, right):
-    return _make_fo_expr(14, left, right)
-
-def eq(left, right):
-    return _make_fo_expr(30, left, right)
-
-def neq(left, right):
-    return _make_fo_expr(31, left, right)
-
-def lte(left, right):
-    return _make_fo_expr(32, left, right)
-
-def gte(left, right):
-    return _make_fo_expr(33, left, right)
-
-def lt(left, right):
-    return _make_fo_expr(34, left, right)
-
-def gt(left, right):
-    return _make_fo_expr(35, left, right)
-
-# First order operators
-
-def abs(arg):
-    return _make_fo_expr(20, arg)
-
-def exp(arg):
-    return _make_fo_expr(21, arg)
-
-def log(arg):
-    return _make_fo_expr(22, arg)
-
-def integer(arg):
-    return _make_fo_expr(41, arg)
-
-def zero_one(arg):
-    return _make_fo_expr(42, arg)
-
-def interval(arg, low, high):
-    return _make_fo_expr(51, low, high, arg)
-
-def integer_interval(arg, low, high):
-    return _make_fo_expr(50, low, high, 1, arg)
-
-def if_then_else(cond, left, right):
-    return _make_fo_expr(60, cond, left, right)
-
-def not_(arg):
-    return _make_fo_expr(61, arg)
-
-def implies(left, right):
-    return _make_fo_expr(62, left, right)
-
-def iff(left, right):
-    return _make_fo_expr(63, left, right)
-
-def xor(left, right):
-    return _make_fo_expr(64, left, right)
-
-def and_(*args):
-    return _make_fo_expr(70, *args)
-
-def or_(*args):
-    return _make_fo_expr(71, *args)
-
-# Aggregate operators
-
-def sum(*args, per=[]) -> Any:
-    return _make_ho_expr(80, args, per)
-
-def product(*args, per=[]) -> Any:
-    return _make_ho_expr(81, args, per)
-
-def min(*args, per=[]) -> Any:
-    return _make_ho_expr(82, args, per)
-
-def max(*args, per=[]) -> Any:
-    return _make_ho_expr(83, args, per)
-
-def count(*args, per=[]) -> Any:
-    return _make_ho_expr(84, args, per)
-
-def all_different(*args, per=[]) -> Any:
-    return _make_ho_expr(90, args, per)
-
-# --------------------------------------------------
-# Symbolic expression helpers
-# --------------------------------------------------
-
-def _make_fo_expr(*args):
-    expr = rel_sv.rel_primitive_solverlib_fo_appl(*args)
-    expr.__class__ = SolverExpression
-    return expr
-
-# TODO(coey) test:
-# dsl.tag(rel_sv.rel_primitive_solverlib_fo_appl, Builtins.Expensive)
-
-def _make_ho_expr(op, args, per):
-    return SolverExpression(dsl.get_graph(), _ho_appl_def, [args, per, [op]])
-
-_ho_appl_def = dsl.build.aggregate_def("_solverlib_ho_appl")
-
-class SolverExpression(dsl.Expression):
-    def __init__(self, graph, op, args):
-        super().__init__(graph, op, args)
-
-def _wrap_expr(e):
-    # If expression is not known to produce a serialized expression string,
-    # wrap it with the identity operation just in case
-    return e if isinstance(e, SolverExpression) else _make_fo_expr(0, e)
-
-# Symbolic expression context, in which some builtin infix operators are redefined
-# TODO(coey) handle comparison chains (e.g. 0 < x < y <= 1) or throw error
-class Operators(dsl.Context):
-    def _supports_binary_op(self, op):
-        return op in _builtin_binary_map
-
-    def _make_binary_op(self, op, left, right):
-        return _make_fo_expr(_builtin_binary_map[op], left, right)
-
-def operators():
-    return Operators(dsl.get_graph())
-
-# Maps for Builtins operator to SolverLib operator ID
-_builtin_binary_map = {
-    Builtins.plus: 10,
-    Builtins.minus: 11,
-    Builtins.mult: 12,
-    Builtins.div: 13,
-    Builtins.pow: 14,
-    Builtins.approx_eq: 30,
-    Builtins.neq: 31,
-    Builtins.lte: 32,
-    Builtins.gte: 33,
-    Builtins.lt: 34,
-    Builtins.gt: 35,
-}
-
-# Concatenate arguments into a string separated by underscores
-def make_string(args:List):
-    string = args[0]
-    for arg in args[1:]:
-        string = rel_sv.concat(rel_sv.concat(string, "_"), arg)
-    return string
-
-
-# --------------------------------------------------
-# Solver
-# --------------------------------------------------
-
-
-@dataclass
-class PollingState:
-    job_id: str
-    continuation_token: str
-    is_done: bool
-    log_to_console: bool
-
-
-class Solver:
-    def __init__(
-        self,
-        solver_name: str,
-        engine_name: str | None = None,
-        engine_size: str | None = None,
-        auto_suspend_mins: int | None = None,
-        resources: SnowflakeResources | None = None,
-    ):
-        self.provider = Provider(resources=resources)
-        self.solver_name = solver_name.lower()
-
-        self.rai_config = self.provider.resources.config
-        settings: dict[str, Any] = {}
-        if "experimental" in self.rai_config:
-            exp_config = self.rai_config.get("experimental", {})
-            if isinstance(exp_config, dict):
-                if "solvers" in exp_config:
-                    settings = exp_config["solvers"].copy()
-
-        # Engine configuration fields are not necessary for the solver
-        # settings so we `pop` them from the settings object. Default
-        # size and auto_suspend_mins are set in the `Provider` methods.
-        engine_name = engine_name or settings.pop("engine", None)
-        if not engine_name:
-            engine_name = self.provider.resources.get_user_based_engine_name()
-        self.engine_name = engine_name
-
-        self.engine_size = engine_size or settings.pop("engine_size", None)
-        self.engine_auto_suspend_mins = auto_suspend_mins or settings.pop("auto_suspend_mins", None)
-
-        # Set default CSV store setting if not already configured
-        if "store" not in settings:
-            settings["store"] = {}
-        if "csv" not in settings["store"]:
-            settings["store"]["csv"] = {}
-        if "enabled" not in settings["store"]["csv"]:
-            settings["store"]["csv"]["enabled"] = True
-
-        # The settings are used when creating a solver engine, they
-        # may configure each individual solver.
-        self.engine_settings = settings
-
-        # Optimistically set the engine object to a `READY` engine to
-        # avoid checking the engine status on each execution.
-        self.engine:Optional[dict[str,Any]] = {"name": engine_name, "state": "READY"}
-
-        return None
-
-    # --------------------------------------------------
-    # Helper
-    # --------------------------------------------------
-    def _auto_create_solver_async(self):
-        name = self.engine_name
-        auto_suspend_mins = self.engine_auto_suspend_mins
-        size = self.engine_size
-        settings = self.engine_settings
-        with Spinner(
-            "Checking solver status",
-            leading_newline=True,
-        ) as spinner:
-            engine = None
-            engines = [e for e in self.provider.list_solvers() if e["name"] == name]
-            assert len(engines) == 1 or len(engines) == 0
-            if len(engines) != 0:
-                engine = engines[0]
-
-            if engine:
-                # TODO (dba) Logic engines support altering the
-                # auto_suspend_mins setting. Currently, we don't have
-                # this capability for solver engines, so users need to
-                # recreate or use another engine. For both the size
-                # and Gurobi configuration the user anyways has to
-                # create a new one as this configuration must happen
-                # when the engine is created.
-                settings_cannot_be_altered_msg = "The configuration of a solver engine happens when the engine is created and _cannot_ be changed. You either need to specify a new engine name or delete the current engine.\n\nSee `solvers.Provider().delete_solver()`."
-                # Make sure that the solver requested is enabled
-                # on the engine.
-                if self.solver_name not in engine["solvers"]:
-                    raise Exception(
-                        f"Solver `{self.solver_name}` is not enabled on `{name}`.\n\n" + settings_cannot_be_altered_msg
-                    )
-
-                # Make sure size and auto_suspend_mins settings match
-                # what the user requests.
-                if size is not None and size != engine["size"]:
-                    raise Exception(
-                        f"Engine `{name}` has size setting of `{engine['size']}` but size is requested to be `{size}`.\n\n" + settings_cannot_be_altered_msg
-                    )
-
-                if auto_suspend_mins is not None and auto_suspend_mins != engine["auto_suspend_mins"]:
-                    raise Exception(
-                        f"Engine `{name}` has auto_suspend_mins setting of `{engine['auto_suspend_mins']}` but auto_suspend_mins is requested to be `{auto_suspend_mins}`.\n\n" + settings_cannot_be_altered_msg
-                    )
-
-                if engine["state"] == "PENDING":
-                    spinner.update_messages(
-                        {
-                            "finished_message": f"Solver {name} is starting",
-                        }
-                    )
-                    pass
-                elif engine["state"] == "SUSPENDED":
-                    spinner.update_messages(
-                        {
-                            "finished_message": f"Resuming solver {name}",
-                        }
-                    )
-                    self.provider.resume_solver_async(name)
-                elif engine["state"] == "READY":
-                    spinner.update_messages(
-                        {
-                            "finished_message": f"Solver {name} is ready",
-                        }
-                    )
-                    pass
-                else:
-                    spinner.update_messages(
-                        {
-                            "message": f"Restarting solver {name}",
-                        }
-                    )
-                    self.provider.delete_solver(name)
-                    engine = None
-            if not engine:
-                # Validate Gurobi config.
-                if self.solver_name == "gurobi":
-                    is_gurobi_configured = False
-                    gurobi_config = settings.get("gurobi", {})
-                    if all(
-                        k in gurobi_config
-                        for k in ["license_secret_name", "external_access_integration"]
-                    ):
-                        is_gurobi_configured = True
-                    if not is_gurobi_configured:
-                        raise Exception(
-                            "Gurobi is not properly configured. You need to provide both `license_secret_name` and `external_access_integration` in its configuration, see https://docs.relational.ai/build/reasoners/prescriptive/solver-backends/gurobi/#usage"
-                        )
-                self.provider.create_solver_async(name, settings=settings, size=size, auto_suspend_mins=auto_suspend_mins)
-                engine = self.provider.get_solver(name)
-                spinner.update_messages(
-                    {
-                        "finished_message": f"Starting solver {name}...",
-                    }
-                )
-
-            self.engine = engine
-
-    def _exec_job(self, payload, log_to_console=True, query_timeout_mins: Optional[int]=None):
-        if self.engine is None:
-            raise Exception("Engine not initialized.")
-
-        with debugging.span("job") as job_span:
-            # Retry logic. If creating a job fails with an engine
-            # related error we will create/resume/... the engine and
-            # retry.
-            try:
-                job_id = self.provider.create_job_async(self.engine["name"], payload, query_timeout_mins=query_timeout_mins)
-            except Exception as e:
-                err_message = str(e).lower()
-                if isinstance(e, ResponseStatusException):
-                    err_message = e.response.json().get("message", "")
-                if any(kw in err_message.lower() for kw in ENGINE_ERRORS + WORKER_ERRORS + ENGINE_NOT_READY_MSGS):
-                    self._auto_create_solver_async()
-                    # Wait until the engine is ready.
-                    poll_with_specified_overhead(lambda: self._is_solver_ready(), 0.1)
-                    job_id = self.provider.create_job_async(self.engine["name"], payload, query_timeout_mins=query_timeout_mins)
-                else:
-                    raise e
-
-            job_span["job_id"] = job_id
-            debugging.event("job_created", job_span, job_id=job_id, engine_name=self.engine["name"], job_type=ENGINE_TYPE_SOLVER)
-            if not isinstance(job_id, str):
-                job_id = ""
-            polling_state = PollingState(job_id, "", False, log_to_console)
-
-            try:
-                with debugging.span("wait", job_id=job_id):
-                    poll_with_specified_overhead(
-                        lambda: self._check_job_status(polling_state), 0.1
-                    )
-            except KeyboardInterrupt as e:
-                print(f"Canceling job {job_id}")
-                self.provider.cancel_job(job_id)
-                raise e
-
-            return job_id
-
-    def _is_solver_ready(self):
-        if self.engine is None:
-            raise Exception("Engine not initialized.")
-
-        result = self.provider.get_solver(self.engine["name"])
-
-        if result is None:
-            raise Exception("No engine available.")
-            
-        self.engine = result
-        state = result["state"]
-        if state != "READY" and state != "PENDING":
-            # Might have suspended or otherwise gone. Recreate.
-            self._auto_create_solver_async()
-        return state == "READY"
-
-    def _check_job_status(self, state):
-        response = self.provider.get_job(state.job_id)
-        assert response, f"No results from get_job('{state.job_id}')"
-
-        status: str = response["state"]
-
-        self._print_solver_logs(state)
-
-        return status == "COMPLETED" or status == "FAILED" or status == "CANCELED"
-
-    def _print_solver_logs(self, state: PollingState):
-        if state.is_done:
-            return
-
-        resp = self.provider.get_job_events(state.job_id, state.continuation_token)
-
-        # Print solver logs to stdout.
-        for event in resp["events"]:
-            if event["type"] == "LogMessage":
-                if state.log_to_console:
-                    print(event["event"]["message"])
-            else:
-                continue
-
-        state.continuation_token = resp["continuation_token"]
-        if state.continuation_token == "":
-            state.is_done = True
-
-
-# --------------------------------------------------
-# Provider
-# --------------------------------------------------
-#
-# TODO (dba) We use an experimental and unified engine API for
-# solvers. Once it is no longer experimental we can remove the
-# provider here and use the normal PyRel provider.
-
-
-class Provider:
-    def __init__(self, resources=None):
-        if not resources:
-            resources = relationalai.Resources()
-        if not isinstance(resources, SnowflakeResources):
-            raise Exception("Solvers are only supported on SPCS.")
-
-        # Type narrowing: resources is confirmed to be SnowflakeResources
-        self.resources: SnowflakeResources = cast(SnowflakeResources, resources)
-        self.direct_access_client: Optional[DirectAccessClient] = None
-
-        if isinstance(self.resources, DirectAccessResources):
-            self.direct_access_client = self.resources.direct_access_client
-
-    def create_solver(
-        self,
-        name: str,
-        size: str | None = None,
-        settings: dict | None = None,
-        auto_suspend_mins: int | None = None,
-    ):
-        if size is None:
-            size = "HIGHMEM_X64_S"
-        if settings is None:
-            settings = {}
-        engine_config: dict[str, Any] = {"settings": settings}
-        if auto_suspend_mins is not None:
-            engine_config["auto_suspend_mins"] = auto_suspend_mins
-        self.resources._exec_sql(
-            f"CALL {APP_NAME}.experimental.create_engine('{ENGINE_TYPE_SOLVER}', '{name}', '{size}', {engine_config});", None
-        )
-
-    def create_solver_async(
-        self,
-        name: str,
-        size: str | None = None,
-        settings: dict | None = None,
-        auto_suspend_mins: int | None = None,
-    ):
-        if size is None:
-            size = "HIGHMEM_X64_S"
-
-        if self.direct_access_client is not None:
-            payload:dict[str, Any] = {
-                "name": name,
-                "settings": settings,
-               }
-            if auto_suspend_mins is not None:
-                payload["auto_suspend_mins"] = auto_suspend_mins
-            if size is not None:
-                payload["size"] = size
-            response = self.direct_access_client.request(
-                "create_engine",
-                payload=payload,
-                path_params={"engine_type": "solver"},
-               )
-            if response.status_code != 200:
-                raise ResponseStatusException(
-                    f"Failed to create engine {name} with size {size}.", response
-                )
-        else:
-            engine_config: dict[str, Any] = {}
-            if settings is not None:
-                engine_config["settings"] = settings
-            if auto_suspend_mins is not None:
-                engine_config["auto_suspend_mins"] = auto_suspend_mins
-            self.resources._exec_sql(
-                f"CALL {APP_NAME}.experimental.create_engine_async('{ENGINE_TYPE_SOLVER}', '{name}', '{size}', {engine_config});",
-                None
-               )
-
-    def delete_solver(self, name: str):
-        if self.direct_access_client is not None:
-            response = self.direct_access_client.request(
-                "delete_engine", path_params = {"engine_type": ENGINE_TYPE_SOLVER, "engine_name": name}
-            )
-            if response.status_code != 200:
-                raise ResponseStatusException("Failed to delete engine.", response)
-            return None
-        else:
-            self.resources._exec_sql(
-                f"CALL {APP_NAME}.experimental.delete_engine('{ENGINE_TYPE_SOLVER}', '{name}');",
-                None
-            )
-
-    def resume_solver_async(self, name: str):
-        if self.direct_access_client is not None:
-            response = self.direct_access_client.request(
-                "resume_engine", path_params = {"engine_type": ENGINE_TYPE_SOLVER, "engine_name": name}
-            )
-            if response.status_code != 200:
-                raise ResponseStatusException("Failed to resume engine.", response)
-            return None
-        else:
-            self.resources._exec_sql(
-                f"CALL {APP_NAME}.experimental.resume_engine_async('{ENGINE_TYPE_SOLVER}', '{name}');",
-                None
-            )
-            return None
-
-    def get_solver(self, name: str):
-        if self.direct_access_client is not None:
-            response = self.direct_access_client.request(
-                "get_engine", path_params = {"engine_type": ENGINE_TYPE_SOLVER, "engine_name": name}
-            )
-            if response.status_code != 200:
-                raise ResponseStatusException("Failed to get engine.", response)
-            solver = response.json()
-            if not solver :
-                return None
-            solver_state = {
-                "name": solver["name"],
-                "id": solver["id"],
-                "size": solver["size"],
-                "state": solver["status"], # callers are expecting 'state'
-                "created_by": solver["created_by"],
-                "created_on": solver["created_on"],
-                "updated_on": solver["updated_on"],
-                "version": solver["version"],
-                "auto_suspend": solver["auto_suspend_mins"],
-                "suspends_at": solver["suspends_at"],
-                "solvers": []
-                if solver["settings"] == ""
-                else [
-                        k
-                        for (k,v) in json.loads(solver["settings"]).items()
-                        if isinstance(v, dict) and v.get("enabled", False)
-                ],
-                }
-            return solver_state
-        else:
-            results = self.resources._exec_sql(
-                f"CALL {APP_NAME}.experimental.get_engine('{ENGINE_TYPE_SOLVER}', '{name}');",
-                None
-               )
-            return solver_list_to_dicts(results)[0]
-
-    def list_solvers(self, state: str | None = None):
-        if self.direct_access_client is not None:
-            response = self.direct_access_client.request(
-                "list_engines"
-            )
-            if response.status_code != 200:
-                raise ResponseStatusException("Failed to list engines.", response)
-            response_content = response.json()
-            if not response_content:
-                return []
-            engines = [
-                {
-                    "name": engine["name"],
-                    "id": engine["id"],
-                    "size": engine["size"],
-                    "state": engine["status"], # callers are expecting 'state'
-                    "created_by": engine["created_by"],
-                    "created_on": engine["created_on"],
-                    "updated_on": engine["updated_on"],
-                    "auto_suspend_mins": engine["auto_suspend_mins"],
-                    "solvers": []
-                    if engine["settings"] == ""
-                    else [
-                        k
-                        for (k, v) in json.loads(engine["settings"]).items()
-                        if isinstance(v, dict) and v.get("enabled", False)
-                           ],
-                }
-                for engine in response_content.get("engines", [])
-                if (state is None or engine.get("status") == state) and (engine.get("type") == ENGINE_TYPE_SOLVER)
-               ]
-            return sorted(engines, key=lambda x: x["name"])
-        else:
-            where_clause = f"WHERE TYPE='{ENGINE_TYPE_SOLVER}'"
-            where_clause = (
-                f"{where_clause} AND STATUS = '{state.upper()}'" if state else where_clause
-               )
-            statement = f"SELECT NAME,ID,SIZE,STATUS,CREATED_BY,CREATED_ON,UPDATED_ON,AUTO_SUSPEND_MINS,SETTINGS FROM {APP_NAME}.experimental.engines {where_clause};"
-            results = self.resources._exec_sql(statement, None)
-            return solver_list_to_dicts(results)
-
-    # --------------------------------------------------
-    # Job API
-    # --------------------------------------------------
-
-    def create_job_async(self, engine_name, payload, query_timeout_mins: Optional[int]=None):
-        payload_json = json.dumps(payload)
-
-        if query_timeout_mins is None and (timeout_value := self.resources.config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS)) is not None:
-            query_timeout_mins = int(timeout_value)
-            
-        if self.direct_access_client is not None:
-            job = {
-                "job_type":ENGINE_TYPE_SOLVER,
-                "worker_name": engine_name,
-                "timeout_mins": query_timeout_mins,
-                "payload": payload_json,
-            }
-            response = self.direct_access_client.request(
-                "create_job",
-                payload=job,
-            )
-            if response.status_code != 200:
-                raise ResponseStatusException("Failed to create job.", response)
-            response_content = response.json()
-            return response_content["id"]
-        else:
-            if query_timeout_mins is not None:
-                sql_string = textwrap.dedent(f"""
-                CALL {APP_NAME}.experimental.exec_job_async('{ENGINE_TYPE_SOLVER}', '{engine_name}', '{payload_json}', null, {query_timeout_mins})
-                """)
-            else:
-                sql_string = textwrap.dedent(f"""
-                CALL {APP_NAME}.experimental.exec_job_async('{ENGINE_TYPE_SOLVER}', '{engine_name}', '{payload_json}')
-                """)
-            res = self.resources._exec_sql(sql_string, None)
-            return res[0]["ID"]
-
-    def list_jobs(self, state=None, limit=None):
-        if self.direct_access_client is not None:
-            response = self.direct_access_client.request(
-                "list_jobs"
-            )
-            if response.status_code != 200:
-                raise ResponseStatusException("Failed to list jobs.", response)
-            response_content = response.json()
-            if not response_content:
-                return []
-            jobs = [
-                {
-                    "id": job["id"],
-                    "state": job["state"],
-                    "created_by": job["created_by"],
-                    "created_on": job["created_on"],
-                    "finished_at": job.get("finished_at", None),
-                    "duration": job["duration"] if "duration" in job else 0,
-                    "solver": json.loads(job["payload"]).get("solver", ""),
-                    "engine": job.get("engine_name", job["worker_name"]),
-                }
-                for job in response_content.get("jobs", [])
-                if state is None or job.get("state") == state
-               ]
-            return sorted(jobs, key=lambda x: x["created_on"], reverse=True)
-        else:
-            state_clause = f"AND STATE = '{state.upper()}'" if state else ""
-            limit_clause = f"LIMIT {limit}" if limit else ""
-            results = self.resources._exec_sql(
-                f"SELECT ID,STATE,CREATED_BY,CREATED_ON,FINISHED_AT,DURATION,PAYLOAD,ENGINE_NAME FROM {APP_NAME}.experimental.jobs where type='{ENGINE_TYPE_SOLVER}' {state_clause} ORDER BY created_on DESC {limit_clause};",
-                None
-               )
-            return job_list_to_dicts(results)
-
-    def get_job(self, id: str):
-        if self.direct_access_client is not None:
-            response = self.direct_access_client.request(
-                "get_job", path_params = {"job_type": ENGINE_TYPE_SOLVER, "job_id": id}
-            )
-            if response.status_code != 200:
-                raise ResponseStatusException("Failed to get job.", response)
-            response_content = response.json()
-            return response_content["job"]
-        else:
-            results = self.resources._exec_sql(
-                f"CALL {APP_NAME}.experimental.get_job('{ENGINE_TYPE_SOLVER}', '{id}');",
-                None
-            )
-            return job_list_to_dicts(results)[0]
-
-    def get_job_events(self, job_id: str, continuation_token: str = ""):
-        if self.direct_access_client is not None:
-            response = self.direct_access_client.request(
-                "get_job_events",
-                path_params = {"job_type": ENGINE_TYPE_SOLVER, "job_id": job_id, "stream_name": "progress"},
-                query_params={"continuation_token": continuation_token},
-            )
-            if response.status_code != 200:
-                raise ResponseStatusException("Failed to get job events.", response)
-            response_content = response.json()
-            if not response_content:
-                return {
-                    "events": [],
-                    "continuation_token": None
-                }
-            return response_content
-        else:
-            results = self.resources._exec_sql(
-                f"SELECT {APP_NAME}.experimental.get_job_events('{ENGINE_TYPE_SOLVER}', '{job_id}', '{continuation_token}');",
-                None
-               )
-            if not results:
-                return {"events": [], "continuation_token": None}
-            row = results[0][0]
-            if not isinstance(row, str):
-                row = ""
-            return json.loads(row)
-
-    def cancel_job(self, id: str):
-        if self.direct_access_client is not None:
-            response = self.direct_access_client.request(
-                "cancel_job", path_params = {"job_type": ENGINE_TYPE_SOLVER, "job_id": id}
-            )
-            if response.status_code != 200:
-                raise ResponseStatusException("Failed to cancel job.", response)
-            return None
-        else:
-            self.resources._exec_sql(
-                f"CALL {APP_NAME}.experimental.cancel_job('{ENGINE_TYPE_SOLVER}', '{id}');",
-                None
-            )
-            return None
-
-
-def solver_list_to_dicts(results):
-    if not results:
-        return []
-    return [
-        {
-            "name": row["NAME"],
-            "id": row["ID"],
-            "size": row["SIZE"],
-            "state": row["STATUS"],  # callers are expecting 'state'
-            "created_by": row["CREATED_BY"],
-            "created_on": row["CREATED_ON"],
-            "updated_on": row["UPDATED_ON"],
-            "auto_suspend_mins": row["AUTO_SUSPEND_MINS"],
-            "solvers": []
-            if row["SETTINGS"] == ""
-            else [
-                k
-                for (k, v) in json.loads(row["SETTINGS"]).items()
-                if isinstance(v, dict) and v.get("enabled", False)
-            ],
-        }
-        for row in results
-    ]
-
-def job_list_to_dicts(results):
-    if not results:
-        return []
-    return [
-        {
-            "id": row["ID"],
-            "state": row["STATE"],
-            "created_by": row["CREATED_BY"],
-            "created_on": row["CREATED_ON"],
-            "finished_at": row["FINISHED_AT"],
-            "duration": row["DURATION"] if "DURATION" in row else 0,
-            "solver": json.loads(row["PAYLOAD"])["solver"],
-            "engine": row["ENGINE_NAME"],
-        }
-        for row in results
-    ]