relationalai 0.13.0.dev0__py3-none-any.whl → 0.13.2__py3-none-any.whl

relationalai/experimental/solvers.py

@@ -0,0 +1,1087 @@
+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)
+
+        # 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
+    ]