Merge pull request #6 from kbase/dev-client

Add SDK method support

Author: MrCreosote

pyproject.toml

@@ -27,6 +27,7 @@ dev = [
     "ipython==9.6.0",
     "pytest==8.4.2",
     "pytest-cov==7.0.0",
+    "requests-mock==1.12.1",
     "semver>=3.0.4",
 ]
 

src/kbase/sdk_baseclient.py

@@ -7,7 +7,10 @@
 import requests as _requests
 import os as _os
 from urllib.parse import urlparse as _urlparse
+import time as _time
+import traceback as _traceback
 from typing import Any
+from urllib3.exceptions import ProtocolError as _ProtocolError
 
 
 # The first version is a pretty basic port from the old baseclient, removing some no longer
@@ -17,12 +20,20 @@
 __version__ = "0.1.0"
 
 
+_EXP_BACKOFF_SEC = [0.1, 0.2, 0.5, 1, 2, 5, 10, 20, 60, 120, 300]
 _CT = "content-type"
 _AJ = "application/json"
 _URL_SCHEME = frozenset(["http", "https"])
 _CHECK_JOB_RETRIES = 3
 
 
+# tested this manually by shortening _EXP_BACKOFF_MS and adding printouts below
+def _get_next_backoff(backoff_index: int = 1) -> tuple[int, float]:
+    if backoff_index < len(_EXP_BACKOFF_SEC) - 1:
+        backoff_index += 1
+    return backoff_index, _EXP_BACKOFF_SEC[backoff_index]
+
+
 class ServerError(Exception):
 
     def __init__(self, name, code, message, data=None, error=None):
@@ -58,18 +69,12 @@ class SDKBaseClient:
         For SDK methods: the url of the callback service.
         For SDK dynamic services: the url of the Service Wizard.
         For other services: the url of the service.
-    timeout - methods will fail if they take longer than this value in seconds.
+    timeout - http requests will fail if they take longer than this value in seconds.
         Default 1800.
     token - a KBase authentication token.
     trust_all_ssl_certificates - set to True to trust self-signed certificates.
         If you don't understand the implications, leave as the default, False.
     lookup_url - set to true when contacting KBase dynamic services.
-    async_job_check_time_ms - the wait time between checking job state for
-        asynchronous jobs run with the run_job method.
-    async_job_check_time_scale_percent - the percentage increase in wait time between async job
-        check attempts.
-    async_job_check_max_time_ms - the maximum time to wait for a job check attempt before
-        failing.
     """
     def __init__(
             self,
@@ -79,9 +84,6 @@ def __init__(
             token: str = None,
             trust_all_ssl_certificates: bool = False,  # Too much of a pain to test
             lookup_url: bool = False,
-            async_job_check_time_ms: int = 100,
-            async_job_check_time_scale_percent: int = 150,
-            async_job_check_max_time_ms: int = 300000
         ):
         if url is None:
             raise ValueError("A url is required")
@@ -93,9 +95,6 @@ def __init__(
         self._headers = {}
         self.trust_all_ssl_certificates = trust_all_ssl_certificates
         self.lookup_url = lookup_url
-        self.async_job_check_time = async_job_check_time_ms / 1000.0
-        self.async_job_check_time_scale_percent = async_job_check_time_scale_percent
-        self.async_job_check_max_time = async_job_check_max_time_ms / 1000.0
         self.token = None
         if token is not None:
             self.token = token
@@ -166,7 +165,48 @@ def _set_up_context(self, service_ver: str = None):
             return {"service_ver": service_ver}
         return None
 
-    def call_method(self, service_method: str, args: list[Any], *, service_ver: str | None = None):
+    def _check_job(self, service: str, job_id: str):
+        return self._call(self.url, service + "._check_job", [job_id])
+
+    def _submit_job(self, service_method: str, args: list[Any], service_ver: str = None):
+        context = self._set_up_context(service_ver)
+        mod, meth = service_method.split(".")
+        return self._call(self.url, mod + "._" + meth + "_submit", args, context)
+
+    def run_job(self, service_method: str, args: list[Any], service_ver: str = None):
+        """
+        Run a SDK method asynchronously.
+        Required arguments:
+        service_method - the service and method to run, e.g. myserv.mymeth.
+        args - a list of arguments to the method.
+        Optional arguments:
+        service_ver - the version of the service to run, e.g. a git hash
+            or dev/beta/release.
+        """
+        mod = service_method.split(".")[0]
+        job_id = self._submit_job(service_method, args, service_ver)
+        backoff_index = -1
+        check_job_failures = 0
+        while check_job_failures < _CHECK_JOB_RETRIES:
+            backoff_index, backoff = _get_next_backoff(backoff_index)
+            _time.sleep(backoff)
+            try:
+                job_state = self._check_job(mod, job_id)
+            except (ConnectionError, _ProtocolError):
+                _traceback.print_exc()
+                check_job_failures += 1
+            else:
+                if job_state["finished"]:
+                    if not job_state["result"]:
+                        return None
+                    if len(job_state["result"]) == 1:
+                        return job_state["result"][0]
+                    return job_state["result"]
+        raise RuntimeError(f"_check_job failed {check_job_failures} times and exceeded limit")
+
+    def call_method(
+        self, service_method: str, args: list[Any], *, service_ver: str | None = None
+    ):
         """
         Call a standard or dynamic service synchronously.
         Required arguments:

test/test_sdk_baseclient.py

@@ -4,16 +4,24 @@
 import os
 import pytest
 import re
+import requests
 from requests.exceptions import HTTPError, ReadTimeout
 import semver
+import shutil
+import subprocess
+import tempfile
 import threading
 import time
+from urllib3.exceptions import ProtocolError
 
 from kbase import sdk_baseclient
 
 
 _VERSION = "0.1.0"
-_MOCKSERVER_PORT = 31590  # should be fine, find an empty port otherwise
+# should be fine, find an empty ports otherwise
+_MOCKSERVER_PORT = 31590
+_CALLBACK_SERVER_PORT = 31591
+_CALLBACK_SERVER_IMAGE = "ghcr.io/kbase/jobrunner:pr-116"
 
 
 @pytest.fixture(scope="module")
@@ -54,6 +62,72 @@ def mockserver():
     server.shutdown()
 
 
+def _wait_for_callback(callback_url: str):
+    interval = 1
+    limit = 120
+    start = time.monotonic()
+    err = None
+    print("waiting for callback server to start")
+    while time.monotonic() - start < limit:
+        try:
+            res = requests.get(callback_url)
+            restext = res.text
+            if res.status_code == 200 and res.text == "[{}]":
+                print(f"Callback server is up at {callback_url}")
+                return
+        except Exception as e:
+            err = e
+        print("waiting for CBS")
+        time.sleep(interval)
+    raise IOError(f"Callback server did not start, last response: {restext}") from err
+
+
+@pytest.fixture(scope="module")
+def callback(url_and_token):
+    # Tried using the temp path pytest fixture but kept getting lots of warnings
+    tmpdir = tempfile.mkdtemp(prefix="callback_server_data_")
+    container_name = f"sdk_baseclient_test_{str(time.time()).replace('.', '_')}"
+    dockercmd = [
+        "docker", "run",
+        "--platform=linux/amd64",  # until we have multiarch images
+        "--name", container_name,
+        "--rm",
+        # TODO SECURITY when CBS allows, use a file instead
+        #               https://github.com/kbase/JobRunner/issues/90
+        "-e", f"KB_AUTH_TOKEN={url_and_token[1]}",
+        "-e", f"KB_BASE_URL={url_and_token[0]}/services/",
+        "-e", f"JOB_DIR={tmpdir}",
+        "-e", "CALLBACK_IP=localhost",
+        "-e", f"CALLBACK_PORT={_CALLBACK_SERVER_PORT}",
+        "-e", "DEBUG_RUNNER=true", # prints logs from containers
+        "-v", "/var/run/docker.sock:/run/docker.sock",
+        "-v", f"{tmpdir}:{tmpdir}",
+        "-p", f"{_CALLBACK_SERVER_PORT}:{_CALLBACK_SERVER_PORT}",
+        _CALLBACK_SERVER_IMAGE
+    ]
+    proc = subprocess.Popen(dockercmd)
+    callback_url = f"http://localhost:{_CALLBACK_SERVER_PORT}"
+    _wait_for_callback(callback_url)
+
+    try:
+        yield callback_url
+    finally:
+        subprocess.check_call(["docker", "stop", container_name])
+        proc.wait(timeout=10)
+        dockercmd = [
+            "docker", "run",
+            "--platform=linux/amd64",  # until we have multiarch images
+            "--name", container_name,
+            "--rm",
+            "-v", f"{tmpdir}:{tmpdir}",
+            "--entrypoint", "bash",
+            _CALLBACK_SERVER_IMAGE,
+            "-c", f"rm -rf {tmpdir}/*",  # need to use bash for the globbing
+        ]
+        subprocess.check_call(dockercmd)
+        shutil.rmtree(tmpdir)
+
+
 def test_version():
     assert sdk_baseclient.__version__ == _VERSION
 
@@ -82,6 +156,7 @@ def test_call_method_basic_passed_token(url_and_token):
 
 
 def test_call_method_basic_env_token(url_and_token):
+    # Tests returning a single value
     os.environ["KB_AUTH_TOKEN"] = url_and_token[1]
     try:
         _test_call_method_basic(url_and_token[0] + "/services/ws", None)
@@ -103,12 +178,11 @@ def _test_call_method_basic(url: str, token: str |  None):
         assert res is None
 
 
-def test_serialize_sets_and_list_return(url_and_token):
-    """
-    Tests
-    * Serializing set and frozenset
-    * Methods that return a list vs. a single value (save_objects).
-    """
+# TODO add test for service that returns > 1 value. Not sure if any services do this
+
+
+def test_serialize_sets(url_and_token):
+    # Tests serializing set and frozenset
     bc = sdk_baseclient.SDKBaseClient(url_and_token[0] + "/services/ws", token=url_and_token[1])
     ws_name = f"sdk_baseclient_test_{time.time()}"
     try:
@@ -231,9 +305,9 @@ def test_dynamic_service(url_and_token):
     ver = res["version"]
     del res["version"]
     assert res == {
-        'git_url': 'https://github.com/kbaseapps/HTMLFileSetServ',
-        'message': '',
-        'state': 'OK',
+        "git_url": "https://github.com/kbaseapps/HTMLFileSetServ",
+        "message": "",
+        "state": "OK",
     }
     assert semver.Version.parse(ver) > semver.Version.parse("0.0.8")
 
@@ -246,8 +320,72 @@ def test_dynamic_service_with_service_version(url_and_token):
     res = bc.call_method("HTMLFileSetServ.status", [], service_ver="0.0.8")
     del res["git_commit_hash"]
     assert res == {
-        'git_url': 'https://github.com/kbaseapps/HTMLFileSetServ',
-        'message': '',
-        'state': 'OK',
+        "git_url": "https://github.com/kbaseapps/HTMLFileSetServ",
+        "message": "",
+        "state": "OK",
         "version": "0.0.8"
     }
+
+
+###
+# Async job tests
+# 
+# All of the 3 ways of calling services use the same underlying _call method, so we don't
+# reiterate those tests every time.
+###
+
+
+def test_run_job_with_service_ver(url_and_token, callback):
+    bc = sdk_baseclient.SDKBaseClient(callback, token=url_and_token[1], timeout=10)
+    res = bc.run_job(
+        "njs_sdk_test_2.run",
+        # force backoff with a wait
+        [{"id": "simplejob2", "wait": 1}],
+        # it seems semvers don't work for unreleased modules
+        service_ver="9d6b868bc0bfdb61c79cf2569ff7b9abffd4c67f"
+    )
+    assert res == {
+        "id": "simplejob2",
+        "name": "njs_sdk_test_2",
+        "hash": "9d6b868bc0bfdb61c79cf2569ff7b9abffd4c67f",
+        "wait": 1,
+    }
+
+
+def test_run_job_no_return(url_and_token, callback):
+    bc = sdk_baseclient.SDKBaseClient(callback, token=url_and_token[1], timeout=10)
+    res = bc.run_job("HelloServiceDeluxe.how_rude", ["Georgette"])
+    assert res is None
+
+
+def test_run_job_list_return(url_and_token, callback):
+    bc = sdk_baseclient.SDKBaseClient(callback, token=url_and_token[1], timeout=10)
+    res = bc.run_job("HelloServiceDeluxe.say_hellos", ["JimBob", "Gengulphus"])
+    assert res == [
+        'Hi JimBob, you santimonious lickspittle', # the dork that wrote this module can't spell
+        'Hi Gengulphus, what a lovely and scintillating person you are',
+    ]
+
+
+def test_run_job_failure(url_and_token, callback, requests_mock):
+    requests_mock.post(callback, [
+        {"json": {"result": ["job_id"]}},
+        {"exc": ConnectionError("oopsie")},
+        {"exc": ProtocolError("oh dang")},
+        {"exc": ConnectionError("so unreliable omg")},
+    ])
+    bc = sdk_baseclient.SDKBaseClient(callback, token=url_and_token[1], timeout=10)
+    with pytest.raises(RuntimeError, match="_check_job failed 3 times and exceeded limit"):
+        bc.run_job("HelloServiceDeluxe.say_hellos", ["JimBob"])
+
+
+def test_run_job_failure_recovery(url_and_token, callback, requests_mock):
+    requests_mock.post(callback, [
+        {"json": {"result": ["job_id"]}},
+        {"exc": ConnectionError("oopsie")},
+        {"exc": ProtocolError("oh dang")},
+        {"json": {"result": [{"finished": 1, "result": ["meh"]}]}},
+    ])
+    bc = sdk_baseclient.SDKBaseClient(callback, token=url_and_token[1], timeout=10)
+    res = bc.run_job("HelloServiceDeluxe.say_hellos", ["JimBob"])
+    assert res == "meh"