diff --git a/clearml/automation/aws_driver.py b/clearml/automation/aws_driver.py index 2bb94cdc..766e05a6 100644 --- a/clearml/automation/aws_driver.py +++ b/clearml/automation/aws_driver.py @@ -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( diff --git a/clearml/automation/cloud_driver.py b/clearml/automation/cloud_driver.py index 146e0537..f2495932 100644 --- a/clearml/automation/cloud_driver.py +++ b/clearml/automation/cloud_driver.py @@ -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(),