def _create_rok_artifact_from_task(self, task): result = task["task"]["result"] snapshot_id = result["event"]["id"] version = result["event"]["version"] obj = result["event"]["object"] bucket = task["task"]["bucket"] artifact_name = task["task"]["action_params"]["params"]["commit_title"] log.info("Creating %s artifact for '%s/%s?version=%s...'", ROK_SNAPSHOT_ARTIFACT_TYPE_NAME, bucket, obj, version) from rok_gw_client.client import RokClient rok_client = RokClient() task_info = rok_client.version_info(bucket, obj, version) members = int(task_info["group_member_count"]) url = task_info["rok_url"] uri = "/rok/buckets/%s/files/%s/versions/%s" % (bucket, obj, version) hash_value = task_info["hash"] property_types = ROK_SNAPSHOT_ARTIFACT_PROPERTIES values = { "name": metadata_store_pb2.Value(string_value=artifact_name), "id": metadata_store_pb2.Value(string_value=snapshot_id), "version": metadata_store_pb2.Value(string_value=version), "object": metadata_store_pb2.Value(string_value=obj), "bucket": metadata_store_pb2.Value(string_value=bucket), "members": metadata_store_pb2.Value(int_value=members), "URL": metadata_store_pb2.Value(string_value=url), "hash": metadata_store_pb2.Value(string_value=hash_value) } custom_properties = dict() for i in range(members): member_name = "member_%s" % i member_obj = task_info.get("group_%s_object" % member_name) member_version = task_info.get("group_%s_version" % member_name) if not member_obj or not member_version: continue member_info = rok_client.version_info(bucket, member_obj, member_version) member_mp = metadata_store_pb2.Value( string_value=member_info.get("meta_mountpoint")) member_url = metadata_store_pb2.Value( string_value=member_info.get("rok_url")) member_hash = metadata_store_pb2.Value( string_value=member_info.get("hash")) custom_properties["%s_URL" % member_name] = member_url custom_properties["%s_mount_point" % member_name] = member_mp custom_properties["%s_hash" % member_name] = member_hash # KFP UI groups Artifacts by run_id/pipeline_name/workspace before # switching to contexts: # https://github.com/kubeflow/pipelines/pull/2852 # https://github.com/kubeflow/pipelines/pull/3485#issuecomment-612722767 custom_properties["run_id"] = metadata_store_pb2.Value( string_value=self.run_uuid) return self._create_artifact_with_type( uri, ROK_SNAPSHOT_ARTIFACT_TYPE_NAME, property_types, values, custom_properties or None)
def snapshot_pipeline_step(pipeline, step, nb_path): """Take a snapshot of a pipeline step with Rok.""" from rok_gw_client.client import RokClient run_uuid = get_run_uuid() bucket = kfp_utils.get_experiment_from_run_id(run_uuid).name obj = "{}-{}".format(pipeline, run_uuid) commit_title = "Step: {}".format(step) commit_message = "Step '{}' of pipeline run '{}'".format(step, run_uuid) environment = json.dumps({ "KALE_PIPELINE_STEP": step, "KALE_NOTEBOOK_PATH": nb_path }) metadata = json.dumps({"environment": environment, "kfp_runid": run_uuid}) params = { "pod": get_pod_name(), "metadata": metadata, "default_container": "main", "commit_title": commit_title, "commit_message": commit_message } rok = RokClient() # Create the bucket in case it does not exist create_rok_bucket(bucket, client=rok) task_info = rok.version_register(bucket, obj, "pod", params, wait=True) print("Successfully created snapshot for step '%s'" % step) print("You can explore the state of the notebook at the beginning" " of this step by spawning a new notebook from the following" " Rok snapshot:") # FIXME: How do we retrieve the base URL of the ROK UI? version = task_info["task"]["result"]["event"]["version"] url_path = ("/rok/buckets/%s/files/%s/versions/%s" % (encode_url_component(bucket), encode_url_component(obj), encode_url_component(version))) print("\n%s\n" % url_path) md_source = ("# Rok autosnapshot\n" "Rok has successfully created a snapshot for step `%s`.\n\n" "To **explore the execution state** at the beginning of " "this step follow the instructions below:\n\n" "1\\. View the [snapshot in the Rok UI](%s).\n\n" "2\\. Copy the Rok URL.\n\n" "3\\. Create a new Notebook Server by using this Rok URL to " "autofill the form." % (step, url_path)) metadata = { "outputs": [{ "storage": "inline", "source": md_source, "type": "markdown" }] } with open("/mlpipeline-ui-metadata.json", "w") as f: json.dump(metadata, f)
def _get_client(): global _client if _client is None: _client = RokClient() return _client
def create_rok_bucket(bucket, client=None): from rok_gw_client.client import RokClient, GatewayClientError if client is None: client = RokClient() # FIXME: Currently the Rok API only supports update-or-create for buckets, # so we do a HEAD first to avoid updating an existing bucket. This # obviously has a small race, which should be removed by extending the Rok # API with an exclusive creation API call. try: return False, client.bucket_info(bucket) except GatewayClientError as e: if e.response.status_code != 404: raise logger.info("Creating bucket: %s", bucket) return client.bucket_create(bucket)
def _get_client(): from rok_gw_client.client import RokClient global _client if _client is None: _client = RokClient() return _client
def get_client(): """Get (init if not exists) the Rok client.""" from rok_gw_client.client import RokClient global _client if _client is None: _client = RokClient() return _client
def create_rok_bucket(bucket, client=None): """Create a new Rok bucket.""" log.info("Creating Rok bucket '%s'...", bucket) from rok_gw_client.client import RokClient, GatewayClientError if client is None: client = RokClient() # FIXME: Currently the Rok API only supports update-or-create for buckets, # so we do a HEAD first to avoid updating an existing bucket. This # obviously has a small race, which should be removed by extending the Rok # API with an exclusive creation API call. try: bucket_info = client.bucket_info(bucket) log.info("Rok bucket '%s' already exists", bucket) return False, bucket_info except GatewayClientError as e: if e.response.status_code != 404: raise created, bucket_info = client.bucket_create(bucket) log.info("Successfully created Rok bucket '%s'", bucket) return created, bucket_info