def _create_inference_job_worker(self, inference_job, trial, replicas):
sub_train_job = self._db.get_sub_train_job(trial.sub_train_job_id)
model = self._db.get_model(sub_train_job.model_id)
service_type = ServiceType.INFERENCE
install_command = parse_model_install_command(model.dependencies,
enable_gpu=False)
environment_vars = {
'POSTGRES_HOST': self._postgres_host,
'POSTGRES_PORT': self._postgres_port,
'POSTGRES_USER': self._postgres_user,
'POSTGRES_DB': self._postgres_db,
'POSTGRES_PASSWORD': self._postgres_password,
'REDIS_HOST': self._redis_host,
'REDIS_PORT': self._redis_port,
'WORKER_INSTALL_COMMAND': install_command,
'CUDA_VISIBLE_DEVICES': '-1' # Hide GPU
}
service = self._create_service(service_type=service_type,
docker_image=model.docker_image,
replicas=replicas,
environment_vars=environment_vars)
self._db.create_inference_job_worker(service_id=service.id,
inference_job_id=inference_job.id,
trial_id=trial.id)
self._db.commit()
return service
def _create_train_job_worker(self, sub_train_job, gpus=1):
model = self._db.get_model(sub_train_job.model_id)
service_type = ServiceType.TRAIN
install_command = parse_model_install_command(model.dependencies,
enable_gpu=(gpus > 0))
environment_vars = {
'WORKER_INSTALL_COMMAND': install_command,
}
service = self._create_service(service_type=service_type,
docker_image=model.docker_image,
environment_vars=environment_vars,
gpus=gpus)
self._db.create_train_job_worker(service_id=service.id,
sub_train_job_id=sub_train_job.id)
self._db.commit()
return service
def _create_inference_job_worker(self, inference_job, trial):
sub_train_job = self._db.get_sub_train_job(trial.sub_train_job_id)
model = self._db.get_model(sub_train_job.model_id)
service_type = ServiceType.INFERENCE
install_command = parse_model_install_command(model.dependencies,
enable_gpu=False)
environment_vars = {
'WORKER_INSTALL_COMMAND': install_command,
}
service = self._create_service(service_type=service_type,
docker_image=model.docker_image,
environment_vars=environment_vars)
self._db.create_inference_job_worker(service_id=service.id,
inference_job_id=inference_job.id,
trial_id=trial.id)
self._db.commit()
return service
def _create_train_job_worker(self, train_job, sub_train_job, replicas):
model = self._db.get_model(sub_train_job.model_id)
service_type = ServiceType.TRAIN
enable_gpu = int(train_job.budget.get(BudgetType.ENABLE_GPU, 0)) > 0
install_command = parse_model_install_command(model.dependencies,
enable_gpu=enable_gpu)
environment_vars = {
'POSTGRES_HOST': self._postgres_host,
'POSTGRES_PORT': self._postgres_port,
'POSTGRES_USER': self._postgres_user,
'POSTGRES_DB': self._postgres_db,
'POSTGRES_PASSWORD': self._postgres_password,
'ADMIN_HOST': self._admin_host,
'ADMIN_PORT': self._admin_port,
'ADVISOR_HOST': self._advisor_host,
'ADVISOR_PORT': self._advisor_port,
'WORKER_INSTALL_COMMAND': install_command,
**({
'CUDA_VISIBLE_DEVICES': -1
} if not enable_gpu else {}) # Hide GPU if not enabled
}
requirements = []
if enable_gpu:
requirements.append(ServiceRequirement.GPU)
service = self._create_service(service_type=service_type,
docker_image=model.docker_image,
replicas=replicas,
environment_vars=environment_vars,
requirements=requirements)
self._db.create_train_job_worker(service_id=service.id,
train_job_id=train_job.id,
sub_train_job_id=sub_train_job.id)
self._db.commit()
return service