use notebook_artifacts_path

kubeflow · Feb 14, 2020 · 09045fb · 09045fb
1 parent 85ba65a
commit 09045fb
Showing 1 changed file with 4 additions and 10 deletions.
diff --git a/xgboost_synthetic/testing/xgboost_test.py b/xgboost_synthetic/testing/xgboost_test.py
@@ -95,16 +95,10 @@ def test_xgboost_synthetic(record_xml_attribute, name, namespace, # pylint: disa
   notebook_artifacts_path = os.path.join(notebook_artifacts_dir, "notebook.html")
   logging.info("Writing notebook artifact to: %s", notebook_artifacts_path)
   os.makedirs(notebook_artifacts_dir, exist_ok=True)
-  # Download notebook html to artifacts
-  with open(notebook_artifacts_path, "w") as f:
-    try:
-      storage_client = storage.Client()
-      bucket = storage_client.get_bucket(nb_bucket)
-      blob = bucket.get_blob(nb_path)
-      blob.download_to_file(f)
-    except Exception as e:
-      logging.info("downloading notebook fail: %s", e)
-      raise e
+  storage_client = storage.Client()
+  bucket = storage_client.get_bucket(nb_bucket)
+  blob = bucket.get_blob(nb_path)
+  blob.download_to_filename(notebook_artifacts_path)
 
   if last_condition.type not in ["Complete"]:
     logging.error("Job didn't complete successfully")