ClearML/clearml
1
0
Fork 0
mirror of https://github.com/clearml/clearml synced 2025-01-31 17:17:00 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add cpu-only support in AWS autoscaler

Browse Source
This commit is contained in:
allegroai 2022-02-26 13:41:23 +02:00
parent 83a71d0fb0
commit f37e036e51
2 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
clearml/automation/aws_driver.py
View File

@ -44,7 +44,7 @@ class AWSDriver(CloudDriver):
def spin_up_worker(self, resource_conf, worker_prefix, queue_name, task_id):
# user_data script will automatically run when the instance is started. it will install the required packages
# for clearml-agent configure it using environment variables and run clearml-agent on the required queue
user_data = self.gen_user_data(worker_prefix, queue_name, task_id)
user_data = self.gen_user_data(worker_prefix, queue_name, task_id, resource_conf.get("cpu_only", False))
ec2 = boto3.client("ec2", **self.creds())
launch_specification = ConfigFactory.from_dict(

4
clearml/automation/cloud_driver.py
View File

@ -118,7 +118,7 @@ class CloudDriver(ABC):
def instance_type_key(self):
"""Return key in configuration for instance type"""
def gen_user_data(self, worker_prefix, queue_name, task_id):
def gen_user_data(self, worker_prefix, queue_name, task_id, cpu_only=False):
return bash_script_template.format(
queue=queue_name,
worker_prefix=worker_prefix,
@ -131,7 +131,7 @@ class CloudDriver(ABC):
secret_key=self.secret_key,
web_server=self.web_server,
bash_script=self.extra_vm_bash_script,
bash_script=("export NVIDIA_VISIBLE_DEVICES=none; " if cpu_only else "") + self.extra_vm_bash_script,
driver_extra=self.driver_bash_extra(task_id),
docker="--docker '{}'".format(self.docker_image) if self.docker_image else "",
instance_id_command=self.instance_id_command(),