Skip to content

Add functionality for creating optionally reserved TPU VMs #43

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
srk97 wants to merge 3 commits into
base: master
Choose a base branch
from
Open

Add functionality for creating optionally reserved TPU VMs #43

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7a5b733
Add functionality for creating optionally reserved TPU VMs
srk97 Oct 14, 2021
262964c
Update gcp.py
srk97 Oct 14, 2021
ef01c9c
Update TPU get
srk97 Oct 14, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 17 additions & 12 deletions cloud/envs/gcp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -305,7 +305,7 @@ def add(self, *args, **kwargs):
return tpu
return super().add(*args, **kwargs)

def get(self, preemptible=True, name=None, version='v3-8', zone=None):
def get(self, mode='preemptible', name=None, version='v3-8', zone=None, tpu_type='tpu-node'):
tpu = None
assert re.match(r"v\d-\d+", version)
for tpu in self.resources:
Expand All @@ -318,20 +318,24 @@ def get(self, preemptible=True, name=None, version='v3-8', zone=None):
break
else:
logger.debug("creating tpu")
tpu = self.up(preemptible=preemptible, name=name, version=version, zone=zone)
tpu = self.up(mode=mode, name=name, version=version, zone=zone, tpu_type=tpu_type)
tpu.in_use()
return tpu

def _up(self, name, ip, preemptible, version, zone, background):
def _up(self, name, ip, mode, version, zone, background, tpu_type):
logger.info("Trying to acquire TPU with name: {} ip: {}".format(name, ip))
cmd = [
"gcloud", "compute", "tpus", "create", name, "--range={}".format(ip),
"--accelerator-type={}".format(version), "--version={}".format(self.tf_version), "--network=default"
]
assert tpu_type in ['tpu-node', 'tpu-vm'], f"TPU type must be one of tpu-node or tpu-vm, got {tpu_type}"
assert mode in ["on_demand", "preemptible", "reserved"], f"TPU mode must be one of on_demand, reserved preemptible - got {mode}"
if tpu_type == 'tpu-vm':
command_insertion = ["gcloud", "alpha", "compute", "tpus", "tpu-vm"]
else:
command_insertion = ["gcloud", "compute", "tpus"]
cmd = command_insertion + ["create", name, "--range={}".format(ip),
"--accelerator-type={}".format(version), "--version={}".format(self.tf_version), "--network=default"]
if zone:
cmd += ["--zone={}".format(zone)]
if preemptible:
cmd += ["--preemptible"]
if mode != "on_demand" :
cmd += [f"--{mode}"]
if background:
cmd += ["--async"]

Expand All @@ -341,17 +345,18 @@ def _up(self, name, ip, preemptible, version, zone, background):

raise Exception("Failed to create TPU with name: {} ip: {} error: \n{}".format(name, ip, err))

def up(self, preemptible=True, background=False, attempts=5, name=None, version='v3-8', zone=None):
def up(self, mode='preemptible', background=False, attempts=5, name=None, version='v3-8', zone=None, tpu_type='tpu-node'):
if not name:
name = self._new_name()
for i in range(attempts):
try:
tpu = self._up(name,
self._new_ip(),
preemptible=preemptible,
mode=mode,
version=version,
zone=zone,
background=background)
background=background
tpu_type=tpu_type)
tpu.manager = self
self.resources.append(tpu)
return tpu
Expand Down