def setup(args):
"""Setup a GKE cluster for TensorFlow jobs.
Args:
args: Command line arguments that control the setup process.
"""
gke = discovery.build("container", "v1")
project = args.project
cluster_name = args.cluster
zone = args.zone
chart = args.chart
machine_type = "n1-standard-8"
cluster_request = {
"cluster": {
"name": cluster_name,
"description": "A GKE cluster for TF.",
"initialNodeCount": 1,
"nodeConfig": {
"machineType":
machine_type,
"oauthScopes": [
"https://www.googleapis.com/auth/cloud-platform",
],
},
# TODO(jlewi): Stop pinning GKE version once 1.8 becomes the default.
"initialClusterVersion": "1.8.1-gke.1",
}
}
if args.accelerators:
# TODO(jlewi): Stop enabling Alpha once GPUs make it out of Alpha
cluster_request["cluster"]["enableKubernetesAlpha"] = True
cluster_request["cluster"]["nodeConfig"]["accelerators"] = []
for accelerator_spec in args.accelerators:
accelerator_type, accelerator_count = accelerator_spec.split(
"=", 1)
cluster_request["cluster"]["nodeConfig"]["accelerators"].append({
"acceleratorCount":
accelerator_count,
"acceleratorType":
accelerator_type,
})
util.create_cluster(gke, project, zone, cluster_request)
util.configure_kubectl(project, zone, cluster_name)
util.load_kube_config()
# Create an API client object to talk to the K8s master.
api_client = k8s_client.ApiClient()
util.setup_cluster(api_client)
if chart.startswith("gs://"):
remote = chart
chart = os.path.join(tempfile.gettempdir(), os.path.basename(chart))
gcs_client = storage.Client(project=project)
bucket_name, path = util.split_gcs_uri(remote)
bucket = gcs_client.get_bucket(bucket_name)
blob = bucket.blob(path)
logging.info("Downloading %s to %s", remote, chart)
blob.download_to_filename(chart)
t = test_util.TestCase()
try:
start = time.time()
util.run([
"helm", "install", chart, "-n", "tf-job", "--wait", "--replace",
"--set", "rbac.install=true,cloud=gke"
])
except subprocess.CalledProcessError as e:
t.failure = "helm install failed;\n" + e.output
finally:
t.time = time.time() - start
t.name = "helm-tfjob-install"
t.class_name = "GKE"
test_util.create_junit_xml_file([t], args.junit_path, gcs_client)
def setup(args):
"""Setup a GKE cluster for TensorFlow jobs.
Args:
args: Command line arguments that control the setup process.
"""
gke = discovery.build("container", "v1")
project = args.project
cluster_name = args.cluster
zone = args.zone
machine_type = "n1-standard-8"
cluster_request = {
"cluster": {
"name": cluster_name,
"description": "A GKE cluster for TF.",
"initialNodeCount": 1,
"nodeConfig": {
"machineType":
machine_type,
"oauthScopes": [
"https://www.googleapis.com/auth/cloud-platform",
],
},
}
}
if args.accelerators:
# TODO(jlewi): Stop enabling Alpha once GPUs make it out of Alpha
cluster_request["cluster"]["enableKubernetesAlpha"] = True
cluster_request["cluster"]["nodeConfig"]["accelerators"] = []
for accelerator_spec in args.accelerators:
accelerator_type, accelerator_count = accelerator_spec.split(
"=", 1)
cluster_request["cluster"]["nodeConfig"]["accelerators"].append({
"acceleratorCount":
accelerator_count,
"acceleratorType":
accelerator_type,
})
util.create_cluster(gke, project, zone, cluster_request)
util.configure_kubectl(project, zone, cluster_name)
util.load_kube_config()
# Create an API client object to talk to the K8s master.
api_client = k8s_client.ApiClient()
t = test_util.TestCase()
try:
start = time.time()
params = {
"tfJobImage": args.image,
"name": "kubeflow-core",
"namespace": args.namespace,
}
component = "core"
account = util.run_and_output(
["gcloud", "config", "get-value", "account", "--quiet"]).strip()
logging.info("Using GCP account %s", account)
util.run([
"kubectl", "create", "clusterrolebinding", "default-admin",
"--clusterrole=cluster-admin", "--user="******"tf-job-operator"
logging.info("Verifying TfJob controller started.")
# TODO(jlewi): We should verify the image of the operator is the correct.
util.wait_for_deployment(api_client, args.namespace,
tf_job_deployment_name)
# Reraise the exception so that the step fails because there's no point
# continuing the test.
except subprocess.CalledProcessError as e:
t.failure = "kubeflow-deploy failed;\n" + (e.output or "")
raise
except util.TimeoutError as e:
t.failure = e.message
raise
finally:
t.time = time.time() - start
t.name = "kubeflow-deploy"
t.class_name = "GKE"
gcs_client = storage.Client(project=args.project)
test_util.create_junit_xml_file([t], args.junit_path, gcs_client)