Make pods wait for user code containers to be actual ready

app/controller/chainsail/controller/__init__.py

@@ -549,10 +549,6 @@ def _do_single_run(self, storage):
         """
         iteration = storage.sim_path
         self._ask_scheduler_to_add_iteration(iteration)
-        # Dirty hack to give nodes time (two mintues) to finish installing packages, compile Stan models etc.
-        # The fact that the controller can just kick off sampling whenever it likes, irrespective of whether the
-        # user code containers are ready, is a bug and tracked in https://github.com/tweag/chainsail/issues/386.
-        time.sleep(120)
         super()._do_single_run(storage)
 
 

app/scheduler/chainsail/scheduler/app.py

@@ -10,7 +10,6 @@
 
 import shortuuid
 from celery import chain
-from cloudstorage.exceptions import NotFoundError
 from firebase_admin.auth import (
     ExpiredIdTokenError,
     InvalidIdTokenError,
@@ -231,7 +230,7 @@ def get_job_signed_url(job_id, user_id):
     find_job(job_id, user_id)
     try:
         signed_url = get_signed_url(job_id)
-    except NotFoundError:
+    except:
         return ("Results not zipped yet", 404)
     update_signed_url_task.apply_async((job_id, signed_url), {})
     return (signed_url, 200)

app/scheduler/chainsail/scheduler/nodes/k8s_pod.py

@@ -215,6 +215,17 @@ def _create_pod(self) -> V1Pod:
                     sub_path=self._CM_FILE_USERCODE,
                 ),
             ],
+            readiness_probe=kub.client.V1Probe(
+                grpc=kub.client.V1GRPCAction(port=50052),
+                # try every 30 seconds...
+                period_seconds=30,
+                # ... up to 10 times.
+                # This gives the user code server 30 * 10 seconds = 5 minutes to be ready.
+                failure_threshold=10,
+                # If everything is well, the gradient evaluation in the health check
+                # should surely not take more than 5 seconds
+                timeout_seconds=5,
+                )
         )
         # Worker container
         container_cmd = [self._config.cmd] + self._config.args

app/scheduler/pyproject.toml

@@ -29,8 +29,8 @@ grpcio = "^1.35.0"
 chainsail-grpc = { path = "../../lib/grpc", develop = true }
 firebase-admin = "^4.5.2"
 uWSGI = "^2.0.20"
-cloudstorage = {extras = ["google"], version = "^0.11.0"}
-kubernetes = "^18.20.0"
+# cloudstorage = {extras = ["google"], version = "^0.11.0"}
+kubernetes = "^26.1.0"
 boto3 = "^1.20.26"
 botocore = "1.31.8"
 setuptools = "*"

app/user_code_server/chainsail/user_code_server/__init__.py

@@ -16,7 +16,6 @@
 logger = logging.getLogger("chainsail.controller")
 pdf, initial_states = import_from_user()
 
-
 class UserCodeServicer(user_code_pb2_grpc.UserCodeServicer):
     def LogProb(self, request, context):
         state = np.frombuffer(request.state_bytes)
@@ -50,6 +49,17 @@ def InitialState(self, request, context):
         return user_code_pb2.InitialStateResponse(initial_state_bytes=initial_states.tobytes())
 
 
+class HealthServicer(user_code_pb2_grpc.HealthServicer):
+        def Check(self, request, context):
+            pdf.log_prob_gradient(initial_states)
+            return user_code_pb2.HealthCheckResponse(status=user_code_pb2.HealthCheckResponse.SERVING)
+
+        def Watch(self, request, context):
+            pdf.log_prob_gradient(initial_states)
+            return user_code_pb2.HealthCheckResponse(status=user_code_pb2.HealthCheckResponse.SERVING)
+
+
+
 @click.command()
 @click.option(
     "--port",
@@ -66,9 +76,10 @@ def InitialState(self, request, context):
 def run(port, remote_logging_config):
     # Configure logging
     configure_logging("chainsail.controller", "DEBUG", remote_logging_config)
-
     logger.debug("Starting user code gRPC server")
-    server = grpc.server(futures.ThreadPoolExecutor(max_workers=1))
+
+    server = grpc.server(futures.ThreadPoolExecutor(max_workers=2))
+    user_code_pb2_grpc.add_HealthServicer_to_server(HealthServicer(), server)
     user_code_pb2_grpc.add_UserCodeServicer_to_server(UserCodeServicer(), server)
     server.add_insecure_port(f"[::]:{port}")
     server.start()

docker/scheduler/Dockerfile

@@ -1,7 +1,7 @@
 ##############################################################################
 # BASE
 ##############################################################################
-FROM python:3.8.7-slim as base
+FROM python:3.9-slim as base
 
 # Set up required directories
 RUN mkdir -p /app/deps && mkdir -p /run/sshd

docker/user-code/Dockerfile

@@ -1,7 +1,7 @@
 ##############################################################################
 # BASE
 ##############################################################################
-FROM python:3.8.7-slim as base
+FROM python:3.9-slim as base
 
 # Set up required directories
 RUN mkdir -p /app/deps && mkdir -p /run/sshd

lib/grpc/chainsail/grpc/health_checking_pb2_grpc.py

@@ -15,15 +15,15 @@ def __init__(self, channel):
             channel: A grpc.Channel.
         """
         self.Check = channel.unary_unary(
-            "/Health/Check",
-            request_serializer=health__checking__pb2.HealthCheckRequest.SerializeToString,
-            response_deserializer=health__checking__pb2.HealthCheckResponse.FromString,
-        )
+                '/Health/Check',
+                request_serializer=health__checking__pb2.HealthCheckRequest.SerializeToString,
+                response_deserializer=health__checking__pb2.HealthCheckResponse.FromString,
+                )
         self.Watch = channel.unary_stream(
-            "/Health/Watch",
-            request_serializer=health__checking__pb2.HealthCheckRequest.SerializeToString,
-            response_deserializer=health__checking__pb2.HealthCheckResponse.FromString,
-        )
+                '/Health/Watch',
+                request_serializer=health__checking__pb2.HealthCheckRequest.SerializeToString,
+                response_deserializer=health__checking__pb2.HealthCheckResponse.FromString,
+                )
 
 
 class HealthServicer(object):
@@ -32,91 +32,68 @@ class HealthServicer(object):
     def Check(self, request, context):
         """Missing associated documentation comment in .proto file."""
         context.set_code(grpc.StatusCode.UNIMPLEMENTED)
-        context.set_details("Method not implemented!")
-        raise NotImplementedError("Method not implemented!")
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
 
     def Watch(self, request, context):
         """Missing associated documentation comment in .proto file."""
         context.set_code(grpc.StatusCode.UNIMPLEMENTED)
-        context.set_details("Method not implemented!")
-        raise NotImplementedError("Method not implemented!")
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
 
 
 def add_HealthServicer_to_server(servicer, server):
     rpc_method_handlers = {
-        "Check": grpc.unary_unary_rpc_method_handler(
-            servicer.Check,
-            request_deserializer=health__checking__pb2.HealthCheckRequest.FromString,
-            response_serializer=health__checking__pb2.HealthCheckResponse.SerializeToString,
-        ),
-        "Watch": grpc.unary_stream_rpc_method_handler(
-            servicer.Watch,
-            request_deserializer=health__checking__pb2.HealthCheckRequest.FromString,
-            response_serializer=health__checking__pb2.HealthCheckResponse.SerializeToString,
-        ),
+            'Check': grpc.unary_unary_rpc_method_handler(
+                    servicer.Check,
+                    request_deserializer=health__checking__pb2.HealthCheckRequest.FromString,
+                    response_serializer=health__checking__pb2.HealthCheckResponse.SerializeToString,
+            ),
+            'Watch': grpc.unary_stream_rpc_method_handler(
+                    servicer.Watch,
+                    request_deserializer=health__checking__pb2.HealthCheckRequest.FromString,
+                    response_serializer=health__checking__pb2.HealthCheckResponse.SerializeToString,
+            ),
     }
-    generic_handler = grpc.method_handlers_generic_handler("Health", rpc_method_handlers)
+    generic_handler = grpc.method_handlers_generic_handler(
+            'Health', rpc_method_handlers)
     server.add_generic_rpc_handlers((generic_handler,))
 
 
-# This class is part of an EXPERIMENTAL API.
+ # This class is part of an EXPERIMENTAL API.
 class Health(object):
     """Missing associated documentation comment in .proto file."""
 
     @staticmethod
-    def Check(
-        request,
-        target,
-        options=(),
-        channel_credentials=None,
-        call_credentials=None,
-        insecure=False,
-        compression=None,
-        wait_for_ready=None,
-        timeout=None,
-        metadata=None,
-    ):
-        return grpc.experimental.unary_unary(
-            request,
+    def Check(request,
             target,
-            "/Health/Check",
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(request, target, '/Health/Check',
             health__checking__pb2.HealthCheckRequest.SerializeToString,
             health__checking__pb2.HealthCheckResponse.FromString,
-            options,
-            channel_credentials,
-            insecure,
-            call_credentials,
-            compression,
-            wait_for_ready,
-            timeout,
-            metadata,
-        )
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
 
     @staticmethod
-    def Watch(
-        request,
-        target,
-        options=(),
-        channel_credentials=None,
-        call_credentials=None,
-        insecure=False,
-        compression=None,
-        wait_for_ready=None,
-        timeout=None,
-        metadata=None,
-    ):
-        return grpc.experimental.unary_stream(
-            request,
+    def Watch(request,
             target,
-            "/Health/Watch",
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_stream(request, target, '/Health/Watch',
             health__checking__pb2.HealthCheckRequest.SerializeToString,
             health__checking__pb2.HealthCheckResponse.FromString,
-            options,
-            channel_credentials,
-            insecure,
-            call_credentials,
-            compression,
-            wait_for_ready,
-            timeout,
-            metadata,
-        )
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)