Support kubernetes runtime for launching jobs

.github/workflows/test.yml

@@ -113,11 +113,16 @@ jobs:
           - search link read kill write mimic workers edit_user sharing_workers
           - resources
           - memoize
-          - copy netcat netcurl
+          - copy
+          - netcat netcurl
           - edit
           - open wopen
           - store_add
         runtime: [docker, kubernetes]
+        exclude:
+          # netcat / netcurl not supported for kubernetes.
+          - test: netcat netcurl
+            runtime: kubernetes
     steps:
       - name: Clear free space
         run: |
@@ -150,14 +155,16 @@ jobs:
           TEST: ${{ matrix.test }}
           VERSION: ${{ github.head_ref || 'master' }}
           CODALAB_LINK_MOUNTS: /tmp
-      - uses: medyagh/setup-minikube@latest
+      - uses: actions/setup-go@v3
         if: matrix.runtime == 'kubernetes'
+        with:
+          go-version: '1.18.1'
       - name: Run tests using Kubernetes runtime
         if: matrix.runtime == 'kubernetes'
         run: |
           sh ./tests/test-setup.sh
           sh ./scripts/local-k8s/setup-ci.sh
-          #python3 test_runner.py --version ${VERSION} ${TEST}
+          python3 test_runner.py --version ${VERSION} ${TEST}
         env:
           TEST: ${{ matrix.test }}
           VERSION: ${{ github.head_ref || 'master' }}
@@ -167,11 +174,16 @@ jobs:
         run: |
           mkdir /tmp/logs
           for c in $(docker ps -a --format="{{.Names}}"); do docker logs $c > /tmp/logs/$c.log 2> /tmp/logs/$c.err.log; done
+      - name: Save kubernetes logs
+        if: always() && matrix.runtime == 'kubernetes'
+        run: |
+          kubectl config use-context kind-codalab
+          kubectl cluster-info dump --output-directory /tmp/logs
       - name: Upload logs
         if: always()
         uses: actions/upload-artifact@v1
         with:
-          name: logs-test-${{ matrix.test }}
+          name: logs-test-${{ matrix.runtime }}-${{ matrix.test }}
           path: /tmp/logs
 
   test_backend_on_worker_restart:

codalab/bin/ws_server.py

@@ -8,6 +8,7 @@
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.WARNING)
+logging.basicConfig(format='%(asctime)s %(message)s %(pathname)s %(lineno)d')
 
 worker_to_ws: Dict[str, Any] = {}
 

codalab/common.py

@@ -44,6 +44,10 @@
 logger.setLevel(logging.WARNING)
 logger = logging.getLogger('apache_beam')
 logger.setLevel(logging.WARNING)
+logger = logging.getLogger('kubernetes')
+logger.setLevel(logging.WARNING)
+logger = logging.getLogger('urllib3')
+logger.setLevel(logging.ERROR)
 
 
 class IntegrityError(ValueError):
@@ -408,4 +412,5 @@ class BundleRuntime(Enum):
     """
 
     DOCKER = "docker"
+    KUBERNETES = "kubernetes"
     SINGULARITY = "singularity"

codalab/worker/main.py

@@ -24,6 +24,8 @@
 from codalab.worker.dependency_manager import DependencyManager
 from codalab.worker.docker_image_manager import DockerImageManager
 from codalab.worker.singularity_image_manager import SingularityImageManager
+from codalab.worker.noop_image_manager import NoOpImageManager
+from codalab.worker.runtime.kubernetes_runtime import KubernetesRuntime
 
 logger = logging.getLogger(__name__)
 
@@ -131,9 +133,13 @@ def parse_args():
     )
     parser.add_argument(
         '--bundle-runtime',
-        choices=[BundleRuntime.DOCKER.value, BundleRuntime.SINGULARITY.value,],
+        choices=[
+            BundleRuntime.DOCKER.value,
+            BundleRuntime.KUBERNETES.value,
+            BundleRuntime.SINGULARITY.value,
+        ],
         default=BundleRuntime.DOCKER.value,
-        help='The runtime through which the worker will run bundles. The options are docker (default) or singularity',
+        help='The runtime through which the worker will run bundles. The options are docker (default), kubernetes, or singularity',
     )
     parser.add_argument(
         '--idle-seconds',
@@ -201,6 +207,21 @@ def parse_args():
     parser.add_argument(
         '--preemptible', action='store_true', help='Whether the worker is preemptible.',
     )
+    parser.add_argument(
+        '--kubernetes-cluster-host',
+        type=str,
+        help='Host address of the Kubernetes cluster. Only applicable if --bundle-runtime is set to kubernetes.',
+    )
+    parser.add_argument(
+        '--kubernetes-auth-token',
+        type=str,
+        help='Kubernetes cluster authorization token. Only applicable if --bundle-runtime is set to kubernetes.',
+    )
+    parser.add_argument(
+        '--kubernetes-cert-path',
+        type=str,
+        help='Path to the SSL cert for the Kubernetes cluster. Only applicable if --bundle-runtime is set to kubernetes.',
+    )
     return parser.parse_args()
 
 
@@ -297,6 +318,15 @@ def main():
         # todo workers with singularity don't work because this is set to none -- handle this
         bundle_runtime_class = None
         docker_runtime = None
+    elif args.bundle_runtime == BundleRuntime.KUBERNETES.value:
+        image_manager = NoOpImageManager()
+        bundle_runtime_class = KubernetesRuntime(
+            args.work_dir,
+            args.kubernetes_auth_token,
+            args.kubernetes_cluster_host,
+            args.kubernetes_cert_path,
+        )
+        docker_runtime = None
     else:
         image_manager = DockerImageManager(
             os.path.join(args.work_dir, 'images-state.json'),
@@ -332,7 +362,7 @@ def main():
         args.checkin_frequency_seconds,
         bundle_service,
         args.shared_file_system,
-        args.tag_exclusive,
+        args.tag_exclusive if args.tag else False,
         args.group,
         ws_server=args.ws_server,
         docker_runtime=docker_runtime,
@@ -342,7 +372,7 @@ def main():
         exit_on_exception=args.exit_on_exception,
         shared_memory_size_gb=args.shared_memory_size_gb,
         preemptible=args.preemptible,
-        bundle_runtime=DockerRuntime(),
+        bundle_runtime=bundle_runtime_class,
     )
 
     # Register a signal handler to ensure safe shutdown.
@@ -394,9 +424,6 @@ def parse_gpuset_args(arg):
     """
     Parse given arg into a set of strings representing gpu UUIDs
     By default, we will try to start a Docker container with nvidia-smi to get the GPUs.
-    If we get an exception that the Docker socket does not exist, which will be the case
-    on Singularity workers, because they do not have root access, and therefore, access to
-    the Docker socket, we should try to get the GPUs with Singularity.
 
     Arguments:
         arg: comma separated string of ints, or "ALL" representing all gpus

codalab/worker/noop_image_manager.py

@@ -0,0 +1,18 @@
+from types import SimpleNamespace
+from codalab.worker.fsm import DependencyStage
+
+
+class NoOpImageManager:
+    """A "no-op" ImageManager. Doesn't do any downloading of images.
+    This is used by the Kubernetes runtime, because Kubernetes itself will take care of image downloading once
+    a pod is launched later.
+    """
+
+    def start(self):
+        pass
+
+    def stop(self):
+        pass
+
+    def get(self, image_spec: str):
+        return SimpleNamespace(stage=DependencyStage.READY, digest=image_spec)

codalab/worker/runtime/__init__.py

@@ -1,11 +1,12 @@
 from typing import Optional, Tuple
 import docker
+from kubernetes.client.rest import ApiException
 
 DEFAULT_RUNTIME = 'runc'  # copied from docker_utils to avoid a circular import
 
 
 # Any errors that relate to runtime API calls failing.
-RuntimeAPIError = (docker.errors.APIError,)
+RuntimeAPIError = (docker.errors.APIError, ApiException)
 
 
 class Runtime: