Add suspend/resume engine to the api

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changelog
 
+## v0.6.20
+
+* Add suspend/resume engine to the api
+* Add examples for suspend/resume
+
 ## v0.6.19
 
 * Fix setting of created_on attribute on AccessToken

examples/resume_engine.py

@@ -0,0 +1,39 @@
+# Copyright 2021 RelationalAI, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+"""Resume an engine."""
+
+from argparse import ArgumentParser
+import json
+import time
+from urllib.request import HTTPError
+from railib import api, config, show
+
+
+def run(engine: str, profile: str):
+    cfg = config.read(profile=profile)
+    ctx = api.Context(**cfg)
+    api.resume_engine_wait(ctx, engine)
+    print(json.dumps(api.get_engine(ctx, engine), indent=2))
+
+
+if __name__ == "__main__":
+    p = ArgumentParser()
+    p.add_argument("engine", type=str, help="engine name")
+    p.add_argument("-p", "--profile", type=str, help="profile name", default="default")
+    args = p.parse_args()
+    try:
+        run(args.engine, args.profile)
+    except HTTPError as e:
+        show.http_error(e)

examples/suspend_engine.py

@@ -0,0 +1,49 @@
+# Copyright 2021 RelationalAI, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+"""Suspend an engine."""
+
+from argparse import ArgumentParser
+import json
+import time
+from urllib.request import HTTPError
+from railib import api, config, show
+
+
+# Answers if the given state is a terminal state.
+def is_term_state(state: str) -> bool:
+    return state == "SUSPENDED" or ("FAILED" in state)
+
+
+def run(engine: str, profile: str):
+    cfg = config.read(profile=profile)
+    ctx = api.Context(**cfg)
+    rsp = api.suspend_engine(ctx, engine)
+    while True:  # wait for request to reach terminal state
+        time.sleep(3)
+        rsp = api.get_engine(ctx, engine)
+        if not rsp or is_term_state(rsp["state"]):
+            break
+    print(json.dumps(rsp, indent=2))
+
+
+if __name__ == "__main__":
+    p = ArgumentParser()
+    p.add_argument("engine", type=str, help="engine name")
+    p.add_argument("-p", "--profile", type=str, help="profile name", default="default")
+    args = p.parse_args()
+    try:
+        run(args.engine, args.profile)
+    except HTTPError as e:
+        show.http_error(e)

railib/__init__.py

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version_info__ = (0, 7, 3)
+__version_info__ = (0, 7, 4)
 __version__ = ".".join(map(str, __version_info__))

railib/api.py

@@ -135,6 +135,8 @@ class Permission(str, Enum):
     "list_oauth_clients",
     "load_csv",
     "update_user",
+    "suspend_engine",
+    "resume_engine",
 ]
 
 
@@ -373,6 +375,30 @@ def create_engine_wait(ctx: Context, engine: str, size: str = "XS", **kwargs):
     return get_engine(ctx, engine)
 
 
+def suspend_engine(ctx: Context, engine: str, **kwargs):
+    data = { "suspend": True }
+    url = _mkurl(ctx, f"{PATH_ENGINE}/{engine}")
+    rsp = rest.patch(ctx, url, data, **kwargs)
+    return json.loads(rsp.read())
+
+
+def resume_engine(ctx: Context, engine: str, **kwargs):
+    data = { "suspend": False }
+    url = _mkurl(ctx, f"{PATH_ENGINE}/{engine}")
+    rsp = rest.patch(ctx, url, data, **kwargs)
+    return json.loads(rsp.read())
+
+
+def resume_engine_wait(ctx: Context, engine: str, **kwargs):
+    resume_engine(ctx, engine, **kwargs)
+    poll_with_specified_overhead(
+        lambda: is_engine_term_state(get_engine(ctx, engine)["state"]),
+        overhead_rate=0.2,
+        timeout=30 * 60,
+    )
+    return get_engine(ctx, engine)
+
+
 def create_user(ctx: Context, email: str, roles: List[Role] = None, **kwargs):
     rs = roles or []
     data = {"email": email, "roles": [r.value for r in rs]}