Support for cloud re-runs and trigger from the UI (#265)

* Recreating the future graph * [MARK] Resolution working, edges not right * RERUN WORKS * Clean-up and refactor * Back to working state * Fixing dag view * mocking storages all around * Linting fixes * pre-commit passes, sending to tests * Aligning mocks * Fix ID label * Fix CloudResolver tests * Cloud tests * Cleanup and documentation * StateMachineResolver docstrings * PR comments * Making methods for testability. * PR comments * UI menus, centralized snackbars * Better ID copy buttons * Refactor top-level state management * Fix clone * Fix submission * UI fixes * Fix edges * Fix nesting bugs * Fixing merge errors * minor fixes * Fix cancelEnabled * last fixes * PR prep * Clone factories tests * Rerun resolution tests. * Mypy fixes * Fix build * Fix build * Apply suggestions from code review Co-authored-by: tscurtu <tudor@sematic.dev> * Apply suggestions from code review Co-authored-by: tscurtu <tudor@sematic.dev> * Cleanup and fixes * PR fixes * Note run id click * Note run id click * Clickable note root id * Fix build * PR comments * BETA chip * Apply suggestions from code review Co-authored-by: tscurtu <tudor@sematic.dev> * Fix RunTime * Revert pkg lock Co-authored-by: tscurtu <tudor@sematic.dev>
sematic-ai · Oct 31, 2022 · 2a11e8d · 2a11e8d
1 parent c050ee3
commit 2a11e8d
Show file tree

Hide file tree

Showing 39 changed files with 1,069 additions and 478 deletions.
diff --git a/sematic/abstract_future.py b/sematic/abstract_future.py
@@ -116,7 +116,7 @@ def __init__(
         retry_settings: Optional[RetrySettings] = None,
         base_image_tag: Optional[str] = None,
     ):
-        self.id: str = uuid.uuid4().hex
+        self.id: str = make_future_id()
         self.calculator = calculator
         self.kwargs = kwargs
         # We don't want to replace futures in kwargs, because it holds
@@ -156,3 +156,7 @@ def __repr__(self):
             f"state={self.state.value}, parent_id={parent_id}, "
             f"nested_id={nested_id}, value={self.value})"
         )
+
+
+def make_future_id() -> str:
+    return uuid.uuid4().hex
diff --git a/sematic/api/BUILD b/sematic/api/BUILD
@@ -1,35 +1,36 @@
 sematic_py_lib(
     name = "app",
     srcs = ["app.py"],
-    deps = [
-        ":wsgi",
-    ],
     pip_deps = [
         "flask",
         "flask-cors",
-    ]
+    ],
+    deps = [
+        ":wsgi",
+    ],
 )
 
 sematic_py_lib(
     name = "server_lib",
     srcs = ["server.py"],
     data = ["//sematic/ui:ui_build"],
+    pip_deps = [
+        "flask",
+        "flask-socketio",
+        "eventlet",
+    ],
     deps = [
         "//sematic:config",
         "//sematic/api:app",
         "//sematic/api/endpoints:artifacts",
         "//sematic/api/endpoints:auth",
         "//sematic/api/endpoints:edges",
+        "//sematic/api/endpoints:events",
         "//sematic/api/endpoints:meta",
         "//sematic/api/endpoints:notes",
         "//sematic/api/endpoints:resolutions",
         "//sematic/api/endpoints:runs",
     ],
-    pip_deps = [
-        "flask",
-        "flask-socketio",
-        "eventlet",
-    ],
 )
 
 py_binary(
@@ -44,11 +45,11 @@ py_binary(
 sematic_py_lib(
     name = "wsgi",
     srcs = ["wsgi.py"],
-    deps = [],
     pip_deps = [
         "flask",
         "gunicorn",
     ],
+    deps = [],
 )
 
 # Does not work just yet on M1 mac

diff --git a/sematic/api/endpoints/BUILD b/sematic/api/endpoints/BUILD
@@ -4,7 +4,6 @@ sematic_py_lib(
     pip_deps = [
         "flask",
         "sqlalchemy",
-        "flask-socketio",
     ],
     # buildifier: leave-alone
     deps = [
@@ -22,6 +21,20 @@ sematic_py_lib(
     ],
 )
 
+sematic_py_lib(
+    name = "events",
+    srcs = ["events.py"],
+    pip_deps = [
+        "flask",
+        "flask-socketio",
+    ],
+    deps = [
+        ":auth",
+        "//sematic/api:app",
+        "//sematic/db/models:user",
+    ],
+)
+
 sematic_py_lib(
     name = "resolutions",
     srcs = ["resolutions.py"],
@@ -32,6 +45,7 @@ sematic_py_lib(
     ],
     deps = [
         ":auth",
+        ":events",
         ":request_parameters",
         "//sematic/api:app",
         "//sematic/db:queries",

diff --git a/sematic/api/endpoints/events.py b/sematic/api/endpoints/events.py
@@ -0,0 +1,51 @@
+# Standard Library
+from typing import Optional
+
+import flask
+
+# Third party
+import flask_socketio  # type: ignore
+
+# Sematic
+from sematic.api.app import sematic_api
+from sematic.api.endpoints.auth import authenticate
+from sematic.db.models.user import User
+
+
+@sematic_api.route("/api/v1/events/<namespace>/<event>", methods=["POST"])
+@authenticate
+def events(user: Optional[User], namespace: str, event: str) -> flask.Response:
+    flask_socketio.emit(
+        event,
+        flask.request.json,
+        namespace="/{}".format(namespace),
+        broadcast=True,
+    )
+    return flask.jsonify({})
+
+
+def broadcast_graph_update(root_id: str) -> None:
+    flask_socketio.emit(
+        "update",
+        dict(run_id=root_id),
+        namespace="/graph",
+        broadcast=True,
+    )
+
+
+def broadcast_resolution_cancel(root_id: str, calculator_path: str) -> None:
+    flask_socketio.emit(
+        "cancel",
+        dict(resolution_id=root_id, calculator_path=calculator_path),
+        namespace="/pipeline",
+        broadcast=True,
+    )
+
+
+def broadcast_pipeline_update(calculator_path: str) -> None:
+    flask_socketio.emit(
+        "update",
+        dict(calculator_path=calculator_path),
+        namespace="/pipeline",
+        broadcast=True,
+    )
diff --git a/sematic/api/endpoints/resolutions.py b/sematic/api/endpoints/resolutions.py
@@ -8,22 +8,28 @@
 
 # Third-party
 import flask
-import flask_socketio  # type: ignore
 import sqlalchemy
 from sqlalchemy.orm.exc import NoResultFound
 
 # Sematic
 from sematic.abstract_future import FutureState
 from sematic.api.app import sematic_api
 from sematic.api.endpoints.auth import authenticate
+from sematic.api.endpoints.events import (
+    broadcast_graph_update,
+    broadcast_pipeline_update,
+    broadcast_resolution_cancel,
+)
 from sematic.api.endpoints.request_parameters import jsonify_error
+from sematic.db.models.factories import clone_resolution, clone_root_run
 from sematic.db.models.resolution import InvalidResolution, Resolution, ResolutionStatus
 from sematic.db.models.run import Run
 from sematic.db.models.user import User
 from sematic.db.queries import (
     get_graph,
     get_resolution,
     get_run,
+    get_run_graph,
     save_graph,
     save_resolution,
 )
@@ -146,6 +152,48 @@ def schedule_resolution_endpoint(
     return flask.jsonify(payload)
 
 
+@sematic_api.route("/api/v1/resolutions/<resolution_id>/rerun", methods=["POST"])
+@authenticate
+def rerun_resolution_endpoint(
+    user: Optional[User], resolution_id: str
+) -> flask.Response:
+    original_resolution = get_resolution(resolution_id)
+
+    if original_resolution.container_image_uri is None:
+        return jsonify_error(
+            (
+                f"Resolution {original_resolution.root_id} cannot be re-run: "
+                "it was initially resolved locally, and therefore "
+                "no container image was built"
+            ),
+            HTTPStatus.BAD_REQUEST,
+        )
+
+    rerun_from = None
+    if flask.request.json and "rerun_from" in flask.request.json:
+        rerun_from = flask.request.json["rerun_from"]
+
+    original_runs, _, original_edges = get_run_graph(original_resolution.root_id)
+    original_root_run = original_runs[0]
+
+    root_run, edges = clone_root_run(original_root_run, original_edges)
+    save_graph(runs=[root_run], edges=edges, artifacts=[])
+
+    resolution = clone_resolution(original_resolution, root_id=root_run.id)
+
+    resolution = schedule_resolution(resolution, rerun_from=rerun_from)
+
+    save_resolution(resolution)
+
+    payload = dict(
+        content=resolution.to_json_encodable(),
+    )
+
+    broadcast_pipeline_update(calculator_path=root_run.calculator_path)
+
+    return flask.jsonify(payload)
+
+
 @sematic_api.route("/api/v1/resolutions/<resolution_id>/cancel", methods=["PUT"])
 @authenticate
 def cancel_resolution_endpoint(
@@ -196,20 +244,9 @@ def cancel_resolution_endpoint(
 
     save_graph(unfinished_runs, [], [])
 
-    flask_socketio.emit(
-        "update",
-        dict(run_id=resolution.root_id),
-        namespace="/graph",
-        broadcast=True,
-    )
-
-    flask_socketio.emit(
-        "cancel",
-        dict(
-            resolution_id=resolution.root_id, calculator_path=root_run.calculator_path
-        ),
-        namespace="/pipeline",
-        broadcast=True,
+    broadcast_graph_update(root_id=resolution.root_id)
+    broadcast_resolution_cancel(
+        root_id=resolution.root_id, calculator_path=root_run.calculator_path
     )
 
     return flask.jsonify(dict(content=resolution.to_json_encodable()))
diff --git a/sematic/api/endpoints/runs.py b/sematic/api/endpoints/runs.py
@@ -13,7 +13,6 @@
 
 # Third-party
 import flask
-import flask_socketio  # type: ignore
 import sqlalchemy
 from sqlalchemy.orm.exc import NoResultFound
 
@@ -383,18 +382,6 @@ def get_run_graph_endpoint(user: Optional[User], run_id: str) -> flask.Response:
     return flask.jsonify(payload)
 
 
-@sematic_api.route("/api/v1/events/<namespace>/<event>", methods=["POST"])
-@authenticate
-def events(user: Optional[User], namespace: str, event: str) -> flask.Response:
-    flask_socketio.emit(
-        event,
-        flask.request.json,
-        namespace="/{}".format(namespace),
-        broadcast=True,
-    )
-    return flask.jsonify({})
-
-
 @sematic_api.route("/api/v1/graph", methods=["PUT"])
 @authenticate
 def save_graph_endpoint(user: Optional[User]):