def connect_to_mlmd() -> metadata_store.MetadataStore: metadata_service_host = os.environ.get( 'METADATA_GRPC_SERVICE_SERVICE_HOST', 'metadata-grpc-service') metadata_service_port = int( os.environ.get('METADATA_GRPC_SERVICE_SERVICE_PORT', 8080)) mlmd_connection_config = metadata_store_pb2.MetadataStoreClientConfig( host=metadata_service_host, port=metadata_service_port, ) # Checking the connection to the Metadata store. for _ in range(100): try: mlmd_store = metadata_store.MetadataStore(mlmd_connection_config) # All get requests fail when the DB is empty, so we have to use a put request. # TODO: Replace with _ = mlmd_store.get_context_types() when https://github.com/google/ml-metadata/issues/28 is fixed _ = mlmd_store.put_execution_type( metadata_store_pb2.ExecutionType(name="DummyExecutionType", )) return mlmd_store except Exception as e: print( 'Failed to access the Metadata store. Exception: "{}"'.format( str(e)), file=sys.stderr) sys.stderr.flush() sleep(1) raise RuntimeError('Could not connect to the Metadata store.')
def _connect(): def establish_connection(store): """Ensure connection to MLMD store by making a request.""" try: _ = store.get_context_types() return True except Exception as e: log.warning( "Failed to access the Metadata store. Exception:" " '%s'", str(e)) return False metadata_service_host = os.environ.get( "METADATA_GRPC_SERVICE_SERVICE_HOST", DEFAULT_METADATA_GRPC_SERVICE_SERVICE_HOST) metadata_service_port = int( os.environ.get("METADATA_GRPC_SERVICE_SERVICE_PORT", DEFAULT_METADATA_GRPC_SERVICE_SERVICE_PORT)) mlmd_connection_config = metadata_store_pb2.MetadataStoreClientConfig( host=metadata_service_host, port=metadata_service_port) mlmd_store = metadata_store.MetadataStore(mlmd_connection_config) # We ensure that the connection to MLMD is established by retrying a # number of times and sleeping for 1 second between the tries. # These numbers are taken from the MetadataWriter implementation. for _ in range(100): if establish_connection(mlmd_store): return mlmd_store time.sleep(1) raise RuntimeError("Could not connect to the Metadata store.")
def __init__(self, grpc_host: str = "metadata-grpc-service.kubeflow", grpc_port: int = 8080, root_certificates: Optional[bytes] = None, private_key: Optional[bytes] = None, certificate_chain: Optional[bytes] = None): """ Args: grpc_host: Required gRPC service host, e.g."metadata-grpc-service.kubeflow". grpc_host: Required gRPC service port. root_certificates: Optional SSL certificate for secure connection. private_key: Optional private_key for secure connection. certificate_chain: Optional certificate_chain for secure connection. The optional parameters are the same as in grpc.ssl_channel_credentials. https://grpc.github.io/grpc/python/grpc.html#grpc.ssl_channel_credentials """ config = mlpb.MetadataStoreClientConfig() config.host = grpc_host config.port = grpc_port if private_key: config.ssl_config.client_key = private_key if root_certificates: config.ssl_config.custom_ca = root_certificates if certificate_chain: config.ssl_config.server_cert = certificate_chain self.store = metadata_store.MetadataStore(config)
def record_pipeline(output_dir: Text, metadata_db_uri: Optional[Text] = None, host: Optional[Text] = None, port: Optional[int] = None, pipeline_name: Optional[Text] = None, run_id: Optional[Text] = None) -> None: """Record pipeline run with run_id to output_dir. For the beam pipeline, metadata_db_uri is required. For KFP pipeline, host and port should be specified. If run_id is not specified, then pipeline_name ought to be specified in order to fetch the latest execution for the specified pipeline. Args: output_dir: Directory path where the pipeline outputs should be recorded. metadata_db_uri: Uri to metadata db. host: Hostname of the metadata grpc server port: Port number of the metadata grpc server. pipeline_name: Pipeline name, which is required if run_id isn't specified. run_id: Pipeline execution run_id. Raises: ValueError: In cases of invalid arguments: - metadata_db_uri is None or host and/or port is None. - run_id is None and pipeline_name is None. FileNotFoundError: if the source artifact uri does not already exist. """ if host is not None and port is not None: metadata_config = metadata_store_pb2.MetadataStoreClientConfig( host=host, port=port) elif metadata_db_uri is not None: metadata_config = metadata.sqlite_metadata_connection_config( metadata_db_uri) else: raise ValueError('For KFP, host and port are required. ' 'For beam pipeline, metadata_db_uri is required.') with metadata.Metadata(metadata_config) as metadata_connection: if run_id is None: if pipeline_name is None: raise ValueError('If the run_id is not specified,' ' pipeline_name should be specified') # fetch executions of the most recently updated execution context. executions = _get_latest_executions(metadata_connection, pipeline_name) else: execution_dict = _get_execution_dict(metadata_connection) if run_id in execution_dict: executions = execution_dict[run_id] else: raise ValueError( 'run_id {} is not recorded in the MLMD metadata'.format( run_id)) for src_uri, dest_uri in _get_paths(metadata_connection, executions, output_dir): io_utils.copy_dir(src_uri, dest_uri) logging.info('Pipeline Recorded at %s', output_dir)
def __init__( self, mlmd_connection_config: Optional[ metadata_store_pb2.MetadataStoreClientConfig] = None, ): if mlmd_connection_config is None: # default to value suitable for local testing mlmd_connection_config = metadata_store_pb2.MetadataStoreClientConfig( host='localhost', port=8080, ) self.mlmd_store = metadata_store.MetadataStore(mlmd_connection_config)
def _get_metadata_store(): if FLAGS.use_grpc_backend: grpc_connection_config = metadata_store_pb2.MetadataStoreClientConfig() if FLAGS.grpc_host is None: raise ValueError("grpc_host argument not set.") grpc_connection_config.host = FLAGS.grpc_host if not FLAGS.grpc_port: raise ValueError("grpc_port argument not set.") grpc_connection_config.port = FLAGS.grpc_port return metadata_store.MetadataStore(grpc_connection_config) connection_config = metadata_store_pb2.ConnectionConfig() connection_config.sqlite.SetInParent() return metadata_store.MetadataStore(connection_config)
def __init__( self, mlmd_connection_config: Optional[ metadata_store_pb2.MetadataStoreClientConfig] = None, ): if mlmd_connection_config is None: # default to value suitable for local testing mlmd_connection_config = metadata_store_pb2.MetadataStoreClientConfig( host='localhost', port=8080, ) self.mlmd_store = metadata_store.MetadataStore(mlmd_connection_config) self.dag_type = self.mlmd_store.get_execution_type( type_name='system.DAGExecution')
def _get_metadata_store(grpc_max_receive_message_length=None): if FLAGS.use_grpc_backend: grpc_connection_config = metadata_store_pb2.MetadataStoreClientConfig() if grpc_max_receive_message_length: (grpc_connection_config.channel_arguments.max_receive_message_length ) = grpc_max_receive_message_length if FLAGS.grpc_host is None: raise ValueError("grpc_host argument not set.") grpc_connection_config.host = FLAGS.grpc_host if not FLAGS.grpc_port: raise ValueError("grpc_port argument not set.") grpc_connection_config.port = FLAGS.grpc_port return metadata_store.MetadataStore(grpc_connection_config) connection_config = metadata_store_pb2.ConnectionConfig() connection_config.sqlite.SetInParent() return metadata_store.MetadataStore(connection_config)
def _get_grpc_metadata_connection_config( kubeflow_metadata_config: kubeflow_pb2.KubeflowGrpcMetadataConfig ) -> metadata_store_pb2.MetadataStoreClientConfig: """Constructs a metadata grpc connection config. Args: kubeflow_metadata_config: Configuration parameters to use for constructing a valid metadata connection config in a Kubeflow cluster. Returns: A metadata_store_pb2.MetadataStoreClientConfig object. """ connection_config = metadata_store_pb2.MetadataStoreClientConfig() connection_config.host = _get_config_value( kubeflow_metadata_config.grpc_service_host) connection_config.port = int( _get_config_value(kubeflow_metadata_config.grpc_service_port)) return connection_config
def main( output_directory: Optional[str] = None, # example host: Optional[str] = None, external_host: Optional[str] = None, launcher_image: Optional[str] = None, experiment: str = 'v2_sample_test_samples', metadata_service_host: Optional[str] = None, metadata_service_port: int = 8080, ): """Test file CLI entrypoint used by Fire. :param host: Hostname pipelines can access, defaults to 'http://ml-pipeline:8888'. :type host: str, optional :param external_host: External hostname users can access from their browsers. :type external_host: str, optional :param output_directory: pipeline output directory that holds intermediate artifacts, example gs://your-bucket/path/to/workdir. :type output_directory: str, optional :param launcher_image: override launcher image, only used in V2_COMPATIBLE mode :type launcher_image: URI, optional :param experiment: experiment the run is added to, defaults to 'v2_sample_test_samples' :type experiment: str, optional :param metadata_service_host: host for metadata grpc service, defaults to METADATA_GRPC_SERVICE_HOST or 'metadata-grpc-service' :type metadata_service_host: str, optional :param metadata_service_port: port for metadata grpc service, defaults to 8080 :type metadata_service_port: int, optional """ # Default to env values, so people can set up their env and run these # tests without specifying any commands. if host is None: host = os.getenv('KFP_HOST', 'http://ml-pipeline:8888') if external_host is None: external_host = host if output_directory is None: output_directory = os.getenv('KFP_OUTPUT_DIRECTORY') if metadata_service_host is None: metadata_service_host = os.getenv('METADATA_GRPC_SERVICE_HOST', 'metadata-grpc-service') if launcher_image is None: launcher_image = os.getenv('KFP_LAUNCHER_IMAGE') client = kfp.Client(host=host) def run_pipeline( pipeline_func: Callable, mode: kfp.dsl.PipelineExecutionMode = kfp.dsl. PipelineExecutionMode.V2_COMPATIBLE, arguments: dict = {}, ) -> kfp_server_api.ApiRunDetail: extra_arguments = {} if mode != kfp.dsl.PipelineExecutionMode.V1_LEGACY: extra_arguments = { kfp.dsl.ROOT_PARAMETER_NAME: output_directory } def _create_run(): return client.create_run_from_pipeline_func( pipeline_func, mode=mode, arguments={ **extra_arguments, **arguments, }, launcher_image=launcher_image, experiment_name=experiment, ) run_result = _retry_with_backoff(fn=_create_run) print("Run details page URL:") print(f"{external_host}/#/runs/details/{run_result.run_id}") run_detail = run_result.wait_for_run_completion(20 * MINUTE) # Hide detailed information for pretty printing workflow_spec = run_detail.run.pipeline_spec.workflow_manifest workflow_manifest = run_detail.pipeline_runtime.workflow_manifest run_detail.run.pipeline_spec.workflow_manifest = None run_detail.pipeline_runtime.workflow_manifest = None pprint(run_detail) # Restore workflow manifest, because test cases may use it run_detail.run.pipeline_spec.workflow_manifest = workflow_spec run_detail.pipeline_runtime.workflow_manifest = workflow_manifest return run_detail # When running locally, port forward MLMD grpc service to localhost:8080 by: # # ```bash # NAMESPACE=kubeflow kubectl port-forward svc/metadata-grpc-service 8080:8080 -n $NAMESPACE # ``` # # Then you can uncomment the following config instead. # mlmd_connection_config = metadata_store_pb2.MetadataStoreClientConfig( # host='localhost', # port=8080, # ) mlmd_connection_config = metadata_store_pb2.MetadataStoreClientConfig( host=metadata_service_host, port=metadata_service_port, ) callback(run_pipeline=run_pipeline, mlmd_connection_config=mlmd_connection_config)
def main( pipeline_root: Optional[str] = None, # example host: Optional[str] = None, external_host: Optional[str] = None, launcher_image: Optional[str] = None, launcher_v2_image: Optional[str] = None, driver_image: Optional[str] = None, experiment: str = 'v2_sample_test_samples', metadata_service_host: Optional[str] = None, metadata_service_port: int = 8080, ): """Test file CLI entrypoint used by Fire. :param host: Hostname pipelines can access, defaults to 'http://ml-pipeline:8888'. :type host: str, optional :param external_host: External hostname users can access from their browsers. :type external_host: str, optional :param pipeline_root: pipeline root that holds intermediate artifacts, example gs://your-bucket/path/to/workdir. :type pipeline_root: str, optional :param launcher_image: override launcher image, only used in V2_COMPATIBLE mode :type launcher_image: URI, optional :param launcher_v2_image: override launcher v2 image, only used in V2_ENGINE mode :type launcher_v2_image: URI, optional :param driver_image: override driver image, only used in V2_ENGINE mode :type driver_image: URI, optional :param experiment: experiment the run is added to, defaults to 'v2_sample_test_samples' :type experiment: str, optional :param metadata_service_host: host for metadata grpc service, defaults to METADATA_GRPC_SERVICE_HOST or 'metadata-grpc-service' :type metadata_service_host: str, optional :param metadata_service_port: port for metadata grpc service, defaults to 8080 :type metadata_service_port: int, optional """ # Default to env values, so people can set up their env and run these # tests without specifying any commands. if host is None: host = os.getenv('KFP_HOST', 'http://ml-pipeline:8888') if external_host is None: external_host = host if pipeline_root is None: pipeline_root = os.getenv('KFP_PIPELINE_ROOT') if not pipeline_root: pipeline_root = os.getenv('KFP_OUTPUT_DIRECTORY') if pipeline_root: logger.warning( f'KFP_OUTPUT_DIRECTORY env var is left for backward compatibility, please use KFP_PIPELINE_ROOT instead.' ) if metadata_service_host is None: metadata_service_host = os.getenv('METADATA_GRPC_SERVICE_HOST', 'metadata-grpc-service') if launcher_image is None: launcher_image = os.getenv('KFP_LAUNCHER_IMAGE') if launcher_v2_image is None: launcher_v2_image = os.getenv('KFP_LAUNCHER_V2_IMAGE') if not launcher_v2_image: raise Exception("launcher_v2_image is empty") if driver_image is None: driver_image = os.getenv('KFP_DRIVER_IMAGE') if not driver_image: raise Exception("driver_image is empty") client = kfp.Client(host=host) def run_pipeline( pipeline_func: Callable, mode: kfp.dsl.PipelineExecutionMode = kfp.dsl. PipelineExecutionMode.V2_COMPATIBLE, enable_caching: bool = False, arguments: Optional[dict] = None, ) -> kfp_server_api.ApiRunDetail: arguments = arguments or {} def _create_run(): if mode == kfp.dsl.PipelineExecutionMode.V2_ENGINE: return run_v2_pipeline( client=client, fn=pipeline_func, driver_image=driver_image, launcher_v2_image=launcher_v2_image, pipeline_root=pipeline_root, enable_caching=enable_caching, arguments={ **arguments, }, ) else: conf = kfp.dsl.PipelineConf() conf.add_op_transformer( # add a default resource request & limit to all container tasks add_default_resource_spec( cpu_request='0.5', cpu_limit='1', memory_limit='512Mi', )) if mode == kfp.dsl.PipelineExecutionMode.V1_LEGACY: conf.add_op_transformer(disable_cache) return client.create_run_from_pipeline_func( pipeline_func, pipeline_conf=conf, mode=mode, arguments=arguments, launcher_image=launcher_image, experiment_name=experiment, # This parameter only works for v2 compatible mode and v2 mode, it does not affect v1 mode enable_caching=enable_caching, ) run_result = _retry_with_backoff(fn=_create_run) print("Run details page URL:") print(f"{external_host}/#/runs/details/{run_result.run_id}") run_detail = run_result.wait_for_run_completion(20 * MINUTE) # Hide detailed information for pretty printing workflow_spec = run_detail.run.pipeline_spec.workflow_manifest workflow_manifest = run_detail.pipeline_runtime.workflow_manifest run_detail.run.pipeline_spec.workflow_manifest = None run_detail.pipeline_runtime.workflow_manifest = None pprint(run_detail) # Restore workflow manifest, because test cases may use it run_detail.run.pipeline_spec.workflow_manifest = workflow_spec run_detail.pipeline_runtime.workflow_manifest = workflow_manifest return run_detail # When running locally, port forward MLMD grpc service to localhost:8080 by: # # ```bash # NAMESPACE=kubeflow kubectl port-forward svc/metadata-grpc-service 8080:8080 -n $NAMESPACE # ``` # # Then you can uncomment the following config instead. # mlmd_connection_config = metadata_store_pb2.MetadataStoreClientConfig( # host='localhost', # port=8080, # ) mlmd_connection_config = metadata_store_pb2.MetadataStoreClientConfig( host=metadata_service_host, port=metadata_service_port, ) callback(run_pipeline=run_pipeline, mlmd_connection_config=mlmd_connection_config)
def main( pipeline_root: Optional[str] = None, # example launcher_v2_image: Optional[str] = None, driver_image: Optional[str] = None, experiment: str = 'v2_sample_test_samples', metadata_service_host: Optional[str] = None, metadata_service_port: int = 8080, ): """Test file CLI entrypoint used by Fire. To configure KFP endpoint, configure env vars following: https://www.kubeflow.org/docs/components/pipelines/sdk/connect- api/#configure-sdk-client-by-environment-variables. KFP UI endpoint can be configured by KF_PIPELINES_UI_ENDPOINT env var. :param pipeline_root: pipeline root that holds intermediate artifacts, example gs://your-bucket/path/to/workdir. :type pipeline_root: str, optional :param launcher_v2_image: override launcher v2 image, only used in V2_ENGINE mode :type launcher_v2_image: URI, optional :param driver_image: override driver image, only used in V2_ENGINE mode :type driver_image: URI, optional :param experiment: experiment the run is added to, defaults to 'v2_sample_test_samples' :type experiment: str, optional :param metadata_service_host: host for metadata grpc service, defaults to METADATA_GRPC_SERVICE_HOST or 'metadata-grpc-service' :type metadata_service_host: str, optional :param metadata_service_port: port for metadata grpc service, defaults to 8080 :type metadata_service_port: int, optional """ if pipeline_root is None: pipeline_root = os.getenv('KFP_PIPELINE_ROOT') if not pipeline_root: pipeline_root = os.getenv('KFP_OUTPUT_DIRECTORY') if pipeline_root: logger.warning( f'KFP_OUTPUT_DIRECTORY env var is left for backward compatibility, please use KFP_PIPELINE_ROOT instead.' ) logger.info(f'KFP_PIPELINE_ROOT={pipeline_root}') if metadata_service_host is None: metadata_service_host = os.getenv('METADATA_GRPC_SERVICE_HOST', 'metadata-grpc-service') logger.info(f'METADATA_GRPC_SERVICE_HOST={metadata_service_host}') if launcher_v2_image is None: launcher_v2_image = os.getenv('KFP_LAUNCHER_V2_IMAGE') if not launcher_v2_image: raise Exception("launcher_v2_image is empty") logger.info(f'KFP_LAUNCHER_V2_IMAGE={launcher_v2_image}') if driver_image is None: driver_image = os.getenv('KFP_DRIVER_IMAGE') if not driver_image: raise Exception("driver_image is empty") logger.info(f'KFP_DRIVER_IMAGE={driver_image}') client = kfp.deprecated.Client() # TODO(Bobgy): avoid using private fields when getting loaded config kfp_endpoint = client._existing_config.host kfp_ui_endpoint = client._uihost logger.info(f'KF_PIPELINES_ENDPOINT={kfp_endpoint}') if kfp_ui_endpoint != kfp_endpoint: logger.info(f'KF_PIPELINES_UI_ENDPOINT={kfp_ui_endpoint}') def run_pipeline( pipeline_func: Optional[Callable], pipeline_file: Optional[str], pipeline_file_compile_path: Optional[str], mode: kfp.deprecated.dsl.PipelineExecutionMode = kfp.deprecated. dsl.PipelineExecutionMode.V2_ENGINE, enable_caching: bool = False, arguments: Optional[dict] = None, dry_run: bool = False, # just compile the pipeline without running it timeout: float = 20 * MINUTE, ) -> kfp_server_api.ApiRunDetail: arguments = arguments or {} def _create_run(): if mode == kfp.deprecated.dsl.PipelineExecutionMode.V2_ENGINE: return run_v2_pipeline( client=client, fn=pipeline_func, file=pipeline_file, driver_image=driver_image, launcher_v2_image=launcher_v2_image, pipeline_root=pipeline_root, enable_caching=enable_caching, arguments={ **arguments, }, ) else: conf = kfp.deprecated.dsl.PipelineConf() conf.add_op_transformer( # add a default resource request & limit to all container tasks add_default_resource_spec( cpu_request='0.5', cpu_limit='1', memory_limit='512Mi', )) if mode == kfp.deprecated.dsl.PipelineExecutionMode.V1_LEGACY: conf.add_op_transformer(_disable_cache) if pipeline_func: return client.create_run_from_pipeline_func( pipeline_func, pipeline_conf=conf, mode=mode, arguments=arguments, experiment_name=experiment, ) else: pyfile = pipeline_file if pipeline_file.endswith(".ipynb"): pyfile = tempfile.mktemp(suffix='.py', prefix="pipeline_py_code") _nb_sample_to_py(pipeline_file, pyfile) if dry_run: subprocess.check_call([sys.executable, pyfile]) return package_path = None if pipeline_file_compile_path: subprocess.check_call([sys.executable, pyfile]) package_path = pipeline_file_compile_path else: package_path = tempfile.mktemp( suffix='.yaml', prefix="kfp_package") from kfp.deprecated.compiler.main import \ compile_pyfile compile_pyfile( pyfile=pyfile, output_path=package_path, mode=mode, pipeline_conf=conf, ) return client.create_run_from_pipeline_package( pipeline_file=package_path, arguments=arguments, experiment_name=experiment, ) run_result = _retry_with_backoff(fn=_create_run) if dry_run: # There is no run_result when dry_run. return print("Run details page URL:") print(f"{kfp_ui_endpoint}/#/runs/details/{run_result.run_id}") run_detail = run_result.wait_for_run_completion(timeout) # Hide detailed information for pretty printing workflow_spec = run_detail.run.pipeline_spec.workflow_manifest workflow_manifest = run_detail.pipeline_runtime.workflow_manifest run_detail.run.pipeline_spec.workflow_manifest = None run_detail.pipeline_runtime.workflow_manifest = None pprint(run_detail) # Restore workflow manifest, because test cases may use it run_detail.run.pipeline_spec.workflow_manifest = workflow_spec run_detail.pipeline_runtime.workflow_manifest = workflow_manifest return run_detail # When running locally, port forward MLMD grpc service to localhost:8080 by: # # 1. NAMESPACE=kubeflow kubectl port-forward svc/metadata-grpc-service 8080:8080 -n $NAMESPACE # 2. Configure env var METADATA_GRPC_SERVICE_HOST=localhost. mlmd_connection_config = metadata_store_pb2.MetadataStoreClientConfig( host=metadata_service_host, port=metadata_service_port, ) callback(run_pipeline=run_pipeline, mlmd_connection_config=mlmd_connection_config)