[batch2] batch2 prototype (#7198)

* [batch2] batch2 prototype

* fix n_jobs change in batch

* fix pipeline lint error

* fix msg in pipeline

* increase timeout for test_batch2

auth/auth/auth.py

@@ -185,15 +185,18 @@ async def rest_logout(request, userdata):
 @routes.get('/api/v1alpha/userinfo')
 async def userinfo(request):
     if 'Authorization' not in request.headers:
+        log.info('Authorization not in request.headers')
         raise web.HTTPUnauthorized()
 
     auth_header = request.headers['Authorization']
     if not auth_header.startswith('Bearer '):
+        log.info('Bearer not in Authorization header')
         raise web.HTTPUnauthorized()
     session_id = auth_header[7:]
 
     # b64 encoding of 32-byte session ID is 44 bytes
     if len(session_id) != 44:
+        log.info('Session id != 44 bytes')
         raise web.HTTPUnauthorized()
 
     dbpool = request.app['dbpool']
@@ -207,6 +210,7 @@ async def userinfo(request):
             users = await cursor.fetchall()
 
     if len(users) != 1:
+        log.info(f'Unknown session id: {session_id}')
         raise web.HTTPUnauthorized()
     user = users[0]
 

auth/create-accounts-pod.yaml

@@ -46,6 +46,15 @@ spec:
            kubectl -n {{ default_ns.name }} create secret generic batch-tokens --from-file=/user-tokens/tokens.json
        fi
        # created batch
+       # create batch2
+       N=$(kubectl -n {{ default_ns.name }} get secret --ignore-not-found=true --no-headers batch2-tokens | wc -l | tr -d '[:space:]')
+       if [[ $N = 0 ]]; then
+           rm -f /user-tokens/tokens.json
+           kubectl -n {{ default_ns.name }} get -o json --export secret test-gsa-key | jq '.metadata.name = "batch-gsa-key"' | kubectl -n {{ default_ns.name }} apply -f -
+           python3 create-session.py '{"username":"batch2","user_id":"hail|batch2","service_account":1,"gsa_email":"hail-test@hail-vdc.iam.gserviceaccount.com","bucket_name":"hail-test-1c9nm","gsa_key_secret_name":"batch-gsa-key","jwt_secret_name":"batch2-tokens"}'
+           kubectl -n {{ default_ns.name }} create secret generic batch2-tokens --from-file=/user-tokens/tokens.json
+       fi
+       # created batch2       
        # create ci
        N=$(kubectl -n {{ default_ns.name }} get secret --ignore-not-found=true --no-headers ci-tokens | wc -l | tr -d '[:space:]')
        if [[ $N = 0 ]]; then

batch/batch/batch.py

@@ -868,7 +868,7 @@ async def from_db_multiple(ids, user):
         return batches
 
     @staticmethod
-    async def create_batch(attributes, callback, userdata):
+    async def create_batch(attributes, callback, userdata, n_jobs):
         user = userdata['username']
 
         id = await db.batch.new_record(
@@ -878,7 +878,8 @@ async def create_batch(attributes, callback, userdata):
             user=user,
             deleted=False,
             cancelled=False,
-            closed=False)
+            closed=False,
+            n_jobs=n_jobs)
 
         batch = Batch(id=id, attributes=attributes, callback=callback,
                       userdata=userdata, user=user, state='running',
@@ -1060,7 +1061,8 @@ async def create_batch(request, userdata):
     batch = await Batch.create_batch(
         attributes=parameters.get('attributes'),
         callback=parameters.get('callback'),
-        userdata=userdata)
+        userdata=userdata,
+        n_jobs=parameters['n_jobs'])
     if batch is None:
         abort(400, f'creation of batch in db failed')
 

batch/batch/schemas.py

@@ -39,5 +39,6 @@
         'keyschema': {'type': 'string'},
         'valueschema': {'type': 'string'}
     },
-    'callback': {'type': 'string'}
+    'callback': {'type': 'string'},
+    'n_jobs': {'type': 'integer', 'nullable': False}
 }

batch/create-batch-tables.sql

@@ -7,7 +7,7 @@ CREATE TABLE IF NOT EXISTS `batch` (
   `deleted` BOOLEAN NOT NULL default false,
   `cancelled` BOOLEAN NOT NULL default false,
   `closed` BOOLEAN NOT NULL default false,
-  `n_jobs` INT NOT NULL default 0,
+  `n_jobs` INT NOT NULL,
   `n_completed` INT NOT NULL default 0,
   `n_succeeded` INT NOT NULL default 0,
   `n_failed` INT NOT NULL default 0,
@@ -59,22 +59,6 @@ CREATE INDEX batch_attributes_key_value ON `batch-attributes` (`key`, `value`(25
 
 DELIMITER $$
 
-CREATE TRIGGER trigger_jobs_insert AFTER INSERT ON jobs
-    FOR EACH ROW BEGIN
-        UPDATE batch SET n_jobs = n_jobs + 1 WHERE id = new.batch_id;
-        IF (NEW.state LIKE 'Error' OR NEW.state LIKE 'Failed' OR NEW.state LIKE 'Success' OR NEW.state LIKE 'Cancelled') THEN
-            UPDATE batch SET n_completed = n_completed + 1 WHERE id = NEW.batch_id;
-            IF (NEW.state LIKE 'Failed' OR NEW.state LIKE 'Error') THEN
-	        UPDATE batch SET n_failed = n_failed + 1 WHERE id = NEW.batch_id;
-            ELSEIF (NEW.state LIKE 'Success') THEN
-                UPDATE batch SET n_succeeded = n_succeeded + 1 WHERE id = NEW.batch_id;
-	    ELSEIF (NEW.state LIKE 'Cancelled') THEN
-                UPDATE batch SET n_cancelled = n_cancelled + 1 WHERE id = NEW.batch_id;
-	    END IF;
-        END IF;
-    END;
-$$
-
 CREATE TRIGGER trigger_jobs_update AFTER UPDATE ON jobs
     FOR EACH ROW BEGIN
         IF (OLD.state NOT LIKE NEW.state) AND (NEW.state LIKE 'Error' OR NEW.state LIKE 'Failed' OR NEW.state LIKE 'Success' OR NEW.state LIKE 'Cancelled') THEN

batch/test/test_aioclient.py

@@ -1,6 +1,4 @@
-import os
 import asyncio
-import aiohttp
 import unittest
 from hailtop.batch_client.aioclient import BatchClient
 

batch2/Dockerfile

@@ -1,10 +1,12 @@
 FROM {{ service_base_image.image }}
 
+COPY batch2/run-worker.sh /
+
 COPY batch2/setup.py batch2/MANIFEST.in /batch/
 COPY batch2/batch /batch/batch/
 RUN pip3 install --no-cache-dir /batch && \
   rm -rf /batch
 
 EXPOSE 5000
 
-CMD ["python3", "-m", "batch.batch"]
+CMD ["python3", "-m", "batch"]

batch2/Dockerfile.check

@@ -0,0 +1,6 @@
+FROM {{ batch_image.image }}
+
+COPY pylintrc pylintrc
+RUN python3 -m pip install --no-cache-dir \
+    pylint==2.3.1 \
+    flake8==3.7.8

batch2/Dockerfile.tables

@@ -0,0 +1,4 @@
+FROM {{ base_image.image }}
+
+COPY create-batch-tables.sql .
+COPY delete-batch-tables.sql .

batch2/Dockerfile.test

@@ -0,0 +1,14 @@
+FROM {{ base_image.image }}
+
+COPY hail/python/setup-hailtop.py /hailtop/setup.py
+COPY hail/python/hailtop /hailtop/hailtop/
+RUN python3 -m pip install --no-cache-dir /hailtop \
+  && rm -rf /hailtop
+
+COPY batch2/test/ /test/
+RUN python3 -m pip install --no-cache-dir \
+  pytest-instafail==0.4.1 \
+  pytest-xdist==1.29.0 \
+  pytest-asyncio==0.10.0
+
+CMD ["python3", "-m", "pytest", "-s", "-v", "--instafail", "-n", "1", "-k", "not test_scale", "/test/"]

batch2/MANIFEST.in

@@ -1,3 +1 @@
 recursive-include batch/templates *
-recursive-include batch/static *
-recursive-include batch/js *

batch2/Makefile

@@ -0,0 +1,38 @@
+PROJECT := $(shell gcloud config get-value project)
+
+BATCH2_LATEST = gcr.io/$(PROJECT)/batch2:latest
+BATCH2_IMAGE = gcr.io/$(PROJECT)/batch2:$(shell docker images -q --no-trunc batch2 | sed -e 's,[^:]*:,,')
+
+PY_FILES = $(shell find batch -iname \*.py -not -exec git check-ignore -q {} \; -print)
+
+PYTHONPATH := $${PYTHONPATH:+$${PYTHONPATH}:}../gear
+PYTHON := PYTHONPATH=$(PYTHONPATH) python3
+
+flake8-stmp: $(PY_FILES)
+	$(PYTHON) -m flake8  --config ../setup.cfg batch
+	touch $@
+
+pylint-stmp: $(PY_FILES)
+	$(PYTHON) -m pylint --rcfile ../pylintrc batch --score=n
+	touch $@
+
+check: flake8-stmp pylint-stmp
+
+build:
+	make -C ../docker build
+	-docker pull $(BATCH2_LATEST)
+	python3 ../ci/jinja2_render.py '{"service_base_image":{"image":"service-base"}}' Dockerfile Dockerfile.out
+	docker build -t batch2 -f Dockerfile.out --cache-from batch2,$(BATCH2_LATEST),service-base ..
+
+push: build
+	docker tag batch2 $(BATCH2_LATEST)
+	docker push $(BATCH2_LATEST)
+	docker tag batch $(BATCH2_IMAGE)
+	docker push $(BATCH2_IMAGE)
+
+deploy: push
+	python3 ../ci/jinja2_render.py '{"code":{"sha":"$(shell git rev-parse --short=12 HEAD)"},"deploy":true,"batch2_image":{"image":"$(BATCH2_IMAGE)"},"batch_pods_ns":{"name":"batch-pods"},"batch_database":{"user_secret_name":"sql-batch-batch-admin-config"},"global":{"domain":"$(DOMAIN)"}}' deployment.yaml deployment.yaml.out
+	kubectl -n default apply -f deployment.yaml.out
+
+clean:
+	rm -f flake8-stmp pylint-stmp

batch2/batch/__main__.py

@@ -0,0 +1,12 @@
+from aiohttp import web
+
+from hailtop.config import get_deploy_config
+from gear import configure_logging
+
+from .batch import app
+
+configure_logging()
+
+deploy_config = get_deploy_config()
+
+web.run_app(deploy_config.prefix_application(app, 'batch2'), host='0.0.0.0', port=5000)