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)
