diff --git a/server/app/routes.py b/server/app/routes.py index a30ffb3a..77e0869a 100644 --- a/server/app/routes.py +++ b/server/app/routes.py @@ -24,7 +24,7 @@ get_cpu_count, get_params, boolean_param, change_parent_task, remove_config_option, model_name_analysis from nmtwizard.capacity import Capacity from nmtwizard.task import TaskEnum, TaskInfos, TasksCreationInfos, TaskPreprocess, TaskTrain, TaskTranslate, \ - TaskScoring, TASK_RELEASE_TYPE + TaskScoring, TaskServe, TASK_RELEASE_TYPE # only for launch() maybe deprecated from nmtwizard.task import TaskBase from utils.storage_utils import StorageUtils @@ -83,7 +83,7 @@ def __init__(self, flask_global, service): 'entity_code': user.entity.entity_code, 'user_code': user.user_code } - self.service_config = config.get_service_config(mongo_client, service_name=GLOBAL_POOL_NAME) + self.service_config = config.get_service_config(mongo_client, GLOBAL_POOL_NAME) self.entity_storages_config = self._get_entity_storages(self.creator['entity_code']) self.storage_client, self.global_storage_name = StorageUtils.get_storages(GLOBAL_POOL_NAME, mongo_client, @@ -849,7 +849,6 @@ def get_translate_score_corpus(testing_data_infos, request_data, routes_config, source = request_data["source"] target = request_data["target"] default_test_data = get_default_test_data(routes_config.storage_client, source, target) if with_default_test else [] - to_translate_corpus = [] to_score_corpus = [] for corpus in testing_data_infos: @@ -1542,7 +1541,7 @@ def launch(service): "registry": get_registry(service_module, image_score), "tag": "latest", "command": ["score", "-o"] + oref["output"] + ["-r"] + oref["ref"] + - option_lang + ['-f', "launcher:scores"] + option_lang + ['-f', "launcher:scores"] } score_task_id, explicit_name = build_task_id(content_score, xxyy, "score", parent_task_id) @@ -1597,7 +1596,7 @@ def launch(service): "registry": get_registry(service_module, image_score), "tag": "latest", "command": ["tuminer", "--tumode", "score", "--srcfile"] + in_out["infile"] + ["--tgtfile"] + - in_out["outfile"] + ["--output"] + in_out["scorefile"] + in_out["outfile"] + ["--output"] + in_out["scorefile"] } tuminer_task_id, explicit_name = build_task_id(content_tuminer, xxyy, "tuminer", parent_task_id) @@ -1800,6 +1799,13 @@ def terminate_internal(task_id): return "problem while posting model: %s" % res, current_status task.terminate(redis_db, task_id, phase=phase) + + # release its port and serve task in mongodb if it is a serve task + if task_id.split("_")[-1] == "serve": + deploy_info = task.get_task_deployment_info(redis_db, task_id) + redis_db.hdel(f"ports:{deploy_info['service']}:{deploy_info['alloc_resource']}", deploy_info['port']) + mongo_client.remove_serving_task_id(deploy_info['model'], task_id) + return "terminating %s" % task_id, current_status diff --git a/server/nmtwizard/common.py b/server/nmtwizard/common.py index f254df08..aca2bf96 100644 --- a/server/nmtwizard/common.py +++ b/server/nmtwizard/common.py @@ -18,7 +18,7 @@ def displaycmd(lst): while p != -1: q = t.find("]]", p) if q != -1: - t = t[0:p] + t[q+2:] + t = t[0:p] + t[q + 2:] else: t = t[0:p] p = t.find("[[private:", p) @@ -45,10 +45,10 @@ def rmprivate(lst): if isinstance(t, six.string_types): p = t.find("[[private:") while p != -1: - t = t[0:p] + t[p+10:] + t = t[0:p] + t[p + 10:] q = t.find("]]") if q != -1: - t = t[0:q] + t[q+2:] + t = t[0:q] + t[q + 2:] p = t.find("[[private:", q) else: p = -1 @@ -165,7 +165,7 @@ def ssh_connect_with_retry(hostname, pkey=pkey, key_filename=key_filename, look_for_keys=False) - logger.info("Connection to %s successful (%f)", hostname, time.time()-start) + logger.info("Connection to %s successful (%f)", hostname, time.time() - start) if login_cmd is not None: if not run_and_check_command(client, login_cmd): raise RuntimeError("failed to run login command") @@ -267,9 +267,9 @@ def cmd_connect_private_registry(docker_registry): if docker_registry['type'] == "aws": return ('$(AWS_ACCESS_KEY_ID=%s AWS_SECRET_ACCESS_KEY=%s ' 'aws ecr get-login --no-include-email --region %s)') % ( - docker_registry['credentials']['AWS_ACCESS_KEY_ID'], - docker_registry['credentials']['AWS_SECRET_ACCESS_KEY'], - docker_registry['region']) + docker_registry['credentials']['AWS_ACCESS_KEY_ID'], + docker_registry['credentials']['AWS_SECRET_ACCESS_KEY'], + docker_registry['region']) username = docker_registry['credentials']['username'] password = docker_registry['credentials']['password'] return 'docker login --username %s --password %s' % (username, password) @@ -291,8 +291,8 @@ def cmd_docker_run(lxpu, docker_options, task_id, if nbgpu == 0 or (nbgpu == 1 and lgpu[0] == 0): gpu_id = '0' else: - env['NV_GPU'] = ",".join([str(int(g)-1) for g in lgpu]) - gpu_id = ",".join([str(v) for v in range(1, nbgpu+1)]) + env['NV_GPU'] = ",".join([str(int(g) - 1) for g in lgpu]) + gpu_id = ",".join([str(v) for v in range(1, nbgpu + 1)]) if docker_options.get('dev') == 1: return "sleep 35" @@ -310,8 +310,16 @@ def cmd_docker_run(lxpu, docker_options, task_id, # launch the task cmd = '%s_o_run_o_-i_o_--rm' % docker_cmd + + # only for serve tasks + if task_id.split("_")[-1] == "serve": + serving_port = docker_command[-1] + cmd += "_o_-p_o_{}:4000".format(serving_port) + docker_command = docker_command[:-1] + if need_expose_gpus: cmd += '_o_--gpus=all' + if 'mount' in docker_options: for k in docker_options['mount']: cmd += '_o_-v_o_%s' % k @@ -416,6 +424,7 @@ def launch_task(task_id, * `callback_url`: server to callback for beat of activity * `callback_interval`: time between 2 beats """ + (lgpu, lcpu) = lxpu gpu_id = ",".join(lgpu) logger.info("launching task - %s / %s", task_id, gpu_id) diff --git a/server/nmtwizard/mongo_database.py b/server/nmtwizard/mongo_database.py index b2984e15..48918e1c 100644 --- a/server/nmtwizard/mongo_database.py +++ b/server/nmtwizard/mongo_database.py @@ -6,7 +6,8 @@ "tags": "pn9-tag", "dockers": "pn9-docker", "evaluations": "pn9-evaluation", - "dataset": "pn9-dataset" + "dataset": "pn9-dataset", + "models": "pn9-model" } @@ -181,3 +182,62 @@ def get_dataset_by_ids(self, dataset_ids): "$in": dataset_ids }}) return dataset + + def create_deployment_info(self, deployment_info): + the_table = self.get("serving_models") + the_table.insert(deployment_info) + + def insert_serving_task(self, model, task_id): + the_table = self.get("models") + query = { + "model": model + } + update = { + "$push": {"tasks": task_id} + } + the_table.update_one(query, update) + + def get_serve_task_id_of_model(self, model): + the_table = self.get("models") + query = { + "model": model + } + model_info = the_table.find_one(query) + tasks = model_info.get("tasks") + for task_id in tasks: + if task_id.split("_")[-1] == "serve": + return task_id + + return None + + def remove_serving_task_id(self, model, task_id): + the_table = self.get("models") + query = { + "model": model + } + + update = { + "$pull": {"tasks": task_id} + } + + the_table.update_one(query, update) + + def get_entity_code_of_model(self, model): + the_table = self.get("models") + query = { + "model": model + } + + return the_table.find_one(query)["owner"]["entity_code"] + + def get_models_of_entity(self, entity_code): + """get all models of the entity""" + the_table = self.get("models") + + query = { + "type": None, + "owner": {"entity_code": entity_code} + } + + return the_table.find(query) + diff --git a/server/nmtwizard/task.py b/server/nmtwizard/task.py index 03d395c5..4ed1a1c0 100644 --- a/server/nmtwizard/task.py +++ b/server/nmtwizard/task.py @@ -248,6 +248,28 @@ def __init__(self, task_infos, parent_task_id, to_translate): TaskBase.__init__(self, task_infos) +class TaskServe(TaskBase): + def __init__(self, task_infos, parent_task_id): + """we set default 4000 for docker serving port, so we set it when creating cmd""" + self._task_suffix = "serve" + self._task_type = "serve" + self._parent_task_id = parent_task_id + + task_infos.content["priority"] = task_infos.content.get("priority", 0) + 1 + task_infos.content["ngpus"] = 0 + if "ncpus" not in task_infos.content: + task_infos.content["ncpus"] = get_cpu_count(task_infos.routes_configuration.service_config, + task_infos.content["ngpus"], "trans") + + task_infos.content["docker"]["command"] = ["-m", self._parent_task_id] + task_infos.content["docker"]["command"].extend(["-ms", "pn9_model_catalog:"]) + task_infos.content["docker"]["command"].extend(["serve"]) + # task_infos.content["docker"]["command"].extend([str(self._serving_port)]) + change_parent_task(task_infos.content["docker"]["command"], parent_task_id) + + TaskBase.__init__(self, task_infos) + + class TaskScoring(TaskBase): def __init__(self, task_infos, parent_task_id, model, to_score): self._task_suffix = "score" @@ -364,7 +386,7 @@ def terminate(redis, task_id, phase): # remove from service queue if it was there service = redis.hget(keyt, "service") if service is not None: - redis.lrem('queued:'+service, 0, task_id) + redis.lrem('queued:' + service, 0, task_id) redis.hset(keyt, "message", phase) set_status(redis, keyt, "terminating") @@ -373,50 +395,66 @@ def terminate(redis, task_id, phase): def work_queue(redis, task_id, service=None, delay=0): if service is None: - service = redis.hget('task:'+task_id, 'service') + service = redis.hget('task:' + task_id, 'service') # Queues the task in the work queue with a delay. if delay == 0: with redis.acquire_lock(f'work_queue:{task_id}', acquire_timeout=1, expire_time=10): if task_id not in redis.lrange(f'work:{service}', 0, -1): - redis.lpush('work:'+service, task_id) - redis.delete('queue:'+task_id) + redis.lpush('work:' + service, task_id) + redis.delete('queue:' + task_id) else: - redis.set('queue:'+task_id, delay) - redis.expire('queue:'+task_id, int(delay)) + redis.set('queue:' + task_id, delay) + redis.expire('queue:' + task_id, int(delay)) def work_unqueue(redis, service): """Pop a task from the work queue.""" - return redis.rpop('work:'+service) + return redis.rpop('work:' + service) def service_queue(redis, task_id, service): """Queue the task on the service queue.""" - with redis.acquire_lock('service:'+service): - redis.lrem('queued:'+service, 0, task_id) - redis.lpush('queued:'+service, task_id) - redis.delete('queue:'+task_id) + with redis.acquire_lock('service:' + service): + redis.lrem('queued:' + service, 0, task_id) + redis.lpush('queued:' + service, task_id) + redis.delete('queue:' + task_id) def enable(redis, task_id, service=None): if service is None: - service = redis.hget('task:'+task_id, 'service') + service = redis.hget('task:' + task_id, 'service') # Marks a task as enabled. - redis.sadd("active:"+service, task_id) + redis.sadd("active:" + service, task_id) def disable(redis, task_id, service=None): if service is None: - service = redis.hget('task:'+task_id, 'service') + service = redis.hget('task:' + task_id, 'service') # Marks a task as disabled. - redis.srem("active:"+service, task_id) - redis.delete("beat:"+task_id) + redis.srem("active:" + service, task_id) + redis.delete("beat:" + task_id) def list_active(redis, service): """Returns all active tasks (i.e. non stopped).""" - return redis.smembers("active:"+service) - + return redis.smembers("active:" + service) + + +def get_task_deployment_info(redis, task_id): + """get deployment info of task: service pool, model, alloc_resource, port""" + task_info = redis.hgetall(f"task:{task_id}") + content = json.loads(task_info["content"]) + model = content["docker"]["command"][1] + alloc_resource = task_info["alloc_resource"] # as trainer name, != host + port = content["docker"]["command"][-1] + service = task_info["service"] + result = { + "service": service, + "model": model, + "alloc_resource": alloc_resource, + "port": port + } + return result def info(redis, taskfile_dir, task_id, fields): """Gets information on a task.""" @@ -459,8 +497,8 @@ def change(redis, task_id, service, priority, ngpus): disable(redis, task_id, prev_service) enable(redis, task_id, service) redis.hset(keyt, "service", service) - redis.lrem('queued:'+prev_service, 0, task_id) - redis.lpush('queued:'+service, task_id) + redis.lrem('queued:' + prev_service, 0, task_id) + redis.lpush('queued:' + service, task_id) if priority: redis.hset(keyt, "priority", priority) if ngpus: @@ -554,7 +592,7 @@ def set_file(redis, taskfile_dir, task_id, content, filename, limit=None): with open(os.path.join(taskdir, filename), "wb") as fh: content = six.ensure_binary(content) if limit and len(content) >= limit: - content = content[:limit-len(disclaimer)] + disclaimer + content = content[:limit - len(disclaimer)] + disclaimer fh.write(content) return content @@ -570,8 +608,8 @@ def append_file(redis, taskfile_dir, task_id, content, filename, limit=None): if limit and current_size >= limit: return content = six.ensure_binary(content, encoding="utf-8") - if limit and len(content)+current_size >= limit: - content = content[:limit-len(disclaimer)] + disclaimer + if limit and len(content) + current_size >= limit: + content = content[:limit - len(disclaimer)] + disclaimer with open(filepath, "ab") as fh: fh.write(content) diff --git a/server/nmtwizard/worker.py b/server/nmtwizard/worker.py index d73381c6..6debaeee 100644 --- a/server/nmtwizard/worker.py +++ b/server/nmtwizard/worker.py @@ -9,6 +9,8 @@ from nmtwizard import task, configuration as config from nmtwizard.capacity import Capacity +from utils.request_utils import RequestUtils +from app import mongo_client def _compatible_resource(resource, request_resource): @@ -172,6 +174,12 @@ def _handle_allocated_task(self, task_id): if v == task_id: lcpu.append(k) self._redis.hset(keyt, 'alloc_lcpu', ",".join(lcpu)) + + if task_id.split('_')[-1] == "serve": + key_port = "ports:%s:%s" % (service.name, resource) + serving_port = content['docker']['command'][-1] + self._redis.hset(key_port, serving_port, task_id) + data = service.launch( task_id, content['options'], @@ -210,11 +218,13 @@ def _handle_allocated_task(self, task_id): task.terminate(self._redis, task_id, phase='launch_error') self._logger.info(traceback.format_exc()) return + self._logger.info('%s: task started on %s', task_id, service.name) self._redis.hset(keyt, 'job', json.dumps(data)) task.set_status(self._redis, keyt, 'running') # For services that do not notify their activity, we should # poll the task status more regularly. + task.work_queue(self._redis, task_id, service.name, delay=service.is_notifying_activity and 120 or 30) @@ -278,6 +288,28 @@ def _block_resource(self, resource, service, err): self._redis.set(keyb, err) self._redis.expire(keyb, self._quarantine_time) + def _select_port_for_deployed_model(self, resource_config): + busy_ports = self._get_all_busy_port_of_resource(self._service, resource_config["name"]) + port_range = resource_config["serve_port_range"] + port_range_from = port_range.get("from") + port_range_to = port_range.get("to") + + if len(busy_ports) == 0: + return port_range_from + + for port in range(port_range_from, port_range_to + 1): + if port not in busy_ports: + return port + return None + + def _get_all_busy_port_of_resource(self, service, resource): + busy_ports = self._redis.hgetall(f"ports:{service}:{resource}") + if busy_ports: + busy_ports = list(map(int, list(busy_ports))) + return busy_ports + else: + return [] + def _allocate_resource(self, task_id, request_resource, service, task_expected_capacity): """Allocates a resource for task_id and returns the name of the resource (or None if none where allocated), and the number of allocated gpus/cpus @@ -668,8 +700,21 @@ def load_machines(self, service_name): request_resource = self._redis.hget(keyt, 'resource') allocated_resource = self._allocate_resource(task_id, request_resource, service, nxpus) if allocated_resource is not None: + # allocate port if it is a serve task + if task_id.split("_")[-1] == "serve": + alloc_resource_config = self._services[self._service]._machines[allocated_resource] + alloc_port = self._select_port_for_deployed_model(alloc_resource_config) + if alloc_port is None: + self._logger.info('%s: there is no port available in resource %s', task_id, allocated_resource) + break + # save port to task content in redis + task_content = json.loads(self._redis.hget(keyt, "content")) + task_content["docker"]["command"].append(alloc_port) + self._redis.hset(keyt, "content", json.dumps(task_content)) + self._logger.info('%s: resource %s reserved %s', task_id, allocated_resource, nxpus) self._redis.hset(keyt, 'alloc_resource', allocated_resource) + task.set_status(self._redis, keyt, 'allocated') task.work_queue(self._redis, task_id, service.name) self._redis.lrem(queue, 0, task_id) diff --git a/server/services/ssh.py b/server/services/ssh.py index b21ae916..b2052482 100644 --- a/server/services/ssh.py +++ b/server/services/ssh.py @@ -149,6 +149,7 @@ def launch(self, support_statistics): options['server'] = resource params = _get_params(self._config, options) + client = self._get_client(params=params) try: callback_url = self._config.get('callback_url') diff --git a/server/utils/request_utils.py b/server/utils/request_utils.py new file mode 100644 index 00000000..c23749e3 --- /dev/null +++ b/server/utils/request_utils.py @@ -0,0 +1,17 @@ +import requests + + +class RequestUtils: + @staticmethod + def get(url, cookies=None): + if cookies is None: + cookies = {} + response = requests.get(url, cookies=cookies) + return response + + @staticmethod + def post(url, json_data=None): + if json_data is None: + json_data = {} + response = requests.post(url, json=json_data) + return response