def _setup_mlmd_port_forward(cls) -> subprocess.Popen:
"""Uses port forward to talk to MLMD gRPC server."""
grpc_port = cls._get_grpc_port()
grpc_forward_command = [
'kubectl', 'port-forward', 'deployment/metadata-deployment', '-n',
'kubeflow', ('%s:%s' % (int(grpc_port) + 1, grpc_port))
]
# Begin port forwarding.
proc = subprocess.Popen(grpc_forward_command)
# Wait while port forward to pod is being established
poll_grpc_port_command = ['lsof', '-i', ':%s' % (int(grpc_port) + 1)]
result = subprocess.run(poll_grpc_port_command) # pylint: disable=subprocess-run-check
timeout = 10
for _ in range(timeout):
if result.returncode == 0:
break
absl.logging.info(
'Waiting while gRPC port-forward is being established...')
time.sleep(1)
result = subprocess.run(poll_grpc_port_command) # pylint: disable=subprocess-run-check
if result.returncode != 0:
raise RuntimeError(
'Failed to establish gRPC port-forward to cluster.')
# Establish MLMD gRPC channel.
forwarding_channel = insecure_channel('localhost:%s' %
(int(grpc_port) + 1))
cls._stub = metadata_store_service_pb2_grpc.MetadataStoreServiceStub(
forwarding_channel)
return proc
def _setup_mlmd_port_forward(cls) -> subprocess.Popen:
"""Uses port forward to talk to MLMD gRPC server."""
grpc_port = cls._get_grpc_port()
is_bind = False
forwarded_port = None
for port in range(_KFP_E2E_TEST_FORWARDING_PORT_BEGIN,
_KFP_E2E_TEST_FORWARDING_PORT_END):
grpc_forward_command = [
'kubectl', 'port-forward',
'deployment/metadata-grpc-deployment', '-n', 'kubeflow',
('%s:%s' % (port, grpc_port))
]
# Begin port forwarding.
proc = subprocess.Popen(grpc_forward_command)
try:
# Wait while port forward to pod is being established
poll_grpc_port_command = ['lsof', '-i', ':%s' % port]
result = subprocess.run( # pylint: disable=subprocess-run-check
poll_grpc_port_command,
stdout=subprocess.PIPE)
for _ in range(_MAX_ATTEMPTS):
if (result.returncode == 0
and 'kubectl' in result.stdout.decode('utf-8')):
is_bind = True
break
logging.info(
'Waiting while gRPC port-forward is being established...'
)
time.sleep(5)
result = subprocess.run( # pylint: disable=subprocess-run-check
poll_grpc_port_command,
stdout=subprocess.PIPE)
except: # pylint: disable=bare-except
# Kill the process in case unexpected error occurred.
proc.kill()
if is_bind:
forwarded_port = port
break
if not is_bind:
raise RuntimeError(
'Failed to establish gRPC port-forward to cluster in '
'the specified range: port %s to %s' %
(_KFP_E2E_TEST_FORWARDING_PORT_BEGIN,
_KFP_E2E_TEST_FORWARDING_PORT_END))
# Establish MLMD gRPC channel.
forwarding_channel = insecure_channel('localhost:%s' % forwarded_port)
cls._stub = metadata_store_service_pb2_grpc.MetadataStoreServiceStub(
forwarding_channel)
return proc
def __init__(self,
config: Union[metadata_store_pb2.ConnectionConfig,
metadata_store_pb2.MetadataStoreClientConfig],
enable_upgrade_migration: bool = False):
"""Initialize the MetadataStore.
MetadataStore can directly connect to either the metadata database or
the metadata store server.
Args:
config: metadata_store_pb2.ConnectionConfig or
metadata_store_pb2.MetadataStoreClientConfig. Configuration to
connect to the database or the metadata store server.
enable_upgrade_migration: if set to True, the library upgrades the db
schema and migrates all data if it connects to an old version backend.
It is ignored when using GRPC client connection config.
"""
self._max_num_retries = 5
if isinstance(config, metadata_store_pb2.ConnectionConfig):
self._using_db_connection = True
migration_options = metadata_store_pb2.MigrationOptions()
migration_options.enable_upgrade_migration = enable_upgrade_migration
self._metadata_store = metadata_store_serialized.CreateMetadataStore(
config.SerializeToString(),
migration_options.SerializeToString())
logging.log(logging.INFO,
'MetadataStore with DB connection initialized')
logging.log(logging.DEBUG, 'ConnectionConfig: %s', config)
if config.HasField('retry_options'):
self._max_num_retries = config.retry_options.max_num_retries
logging.log(
logging.INFO,
'retry options is overwritten: max_num_retries = %d',
self._max_num_retries)
return
if not isinstance(config,
metadata_store_pb2.MetadataStoreClientConfig):
raise ValueError('MetadataStore is expecting either '
'metadata_store_pb2.ConnectionConfig or '
'metadata_store_pb2.MetadataStoreClientConfig')
self._grpc_timeout_sec = None
self._using_db_connection = False
if enable_upgrade_migration:
raise ValueError(
'Upgrade migration is not allowed when using gRPC '
'connection client. Upgrade needs to be performed on '
'the server side.')
channel = self._get_channel(config)
self._metadata_store_stub = (
metadata_store_service_pb2_grpc.MetadataStoreServiceStub(channel))
logging.log(logging.INFO,
'MetadataStore with gRPC connection initialized')
logging.log(logging.DEBUG, 'ConnectionConfig: %s', config)
def __init__(self,
config: Union[metadata_store_pb2.ConnectionConfig,
metadata_store_pb2.MetadataStoreClientConfig],
enable_upgrade_migration: bool = False):
"""Initialize the MetadataStore.
MetadataStore can directly connect to either the metadata database or
the metadata store server.
Args:
config: metadata_store_pb2.ConnectionConfig or
metadata_store_pb2.MetadataStoreClientConfig. Configuration to
connect to the database or the metadata store server.
enable_upgrade_migration: if set to True, the library upgrades the db
schema and migrates all data if it connects to an old version backend.
It is ignored when using GRPC client connection config.
"""
if isinstance(config, metadata_store_pb2.ConnectionConfig):
self._using_db_connection = True
migration_options = metadata_store_pb2.MigrationOptions()
migration_options.enable_upgrade_migration = enable_upgrade_migration
self._metadata_store = metadata_store_serialized.CreateMetadataStore(
config.SerializeToString(),
migration_options.SerializeToString())
logging.log(logging.INFO,
'MetadataStore with DB connection initialized')
return
if not isinstance(config,
metadata_store_pb2.MetadataStoreClientConfig):
raise ValueError('MetadataStore is expecting either'
'metadata_store_pb2.ConnectionConfig or'
'metadata_store_pb2.MetadataStoreClientConfig')
self._using_db_connection = False
if enable_upgrade_migration:
raise ValueError(
'Upgrade migration is not allowed when using gRPC '
'connection client. Upgrade needs to be performed on '
'the server side.')
target = ':'.join([config.host, str(config.port)])
channel = self._get_channel(config, target)
self._metadata_store_stub = (
metadata_store_service_pb2_grpc.MetadataStoreServiceStub(channel))
logging.log(logging.INFO,
'MetadataStore with gRPC connection initialized')
def _setup_mlmd_port_forward(cls) -> subprocess.Popen:
"""Uses port forward to talk to MLMD gRPC server."""
grpc_port = cls._get_grpc_port()
grpc_forward_command = [
'kubectl', 'port-forward', 'deployment/metadata-grpc-deployment',
'-n', 'kubeflow',
('%s:%s' % (_KFP_E2E_TEST_FORWARDING_PORT, grpc_port))
]
# Begin port forwarding.
proc = subprocess.Popen(grpc_forward_command)
try:
# Wait while port forward to pod is being established
poll_grpc_port_command = [
'lsof', '-i',
':%s' % _KFP_E2E_TEST_FORWARDING_PORT
]
result = subprocess.run(poll_grpc_port_command) # pylint: disable=subprocess-run-check
max_attempts = 30
for _ in range(max_attempts):
if result.returncode == 0:
break
logging.info(
'Waiting while gRPC port-forward is being established...')
time.sleep(5)
result = subprocess.run(poll_grpc_port_command) # pylint: disable=subprocess-run-check
except:
# Kill the process in case unexpected error occurred.
proc.kill()
raise
if result.returncode != 0:
raise RuntimeError(
'Failed to establish gRPC port-forward to cluster after %s retries. '
'See error message: %s' % (max_attempts, result.stderr))
# Establish MLMD gRPC channel.
forwarding_channel = insecure_channel('localhost:%s' %
(int(grpc_port) + 1))
cls._stub = metadata_store_service_pb2_grpc.MetadataStoreServiceStub(
forwarding_channel)
return proc
def __init__(self,
config: Union[metadata_store_pb2.ConnectionConfig,
metadata_store_pb2.MetadataStoreClientConfig],
disable_upgrade_migration: bool = False):
"""Initialize the MetadataStore.
MetadataStore can directly connect to either the metadata database or
the metadata store server.
Args:
config: metadata_store_pb2.ConnectionConfig or
metadata_store_pb2.MetadataStoreClientConfig. Configuration to
connect to the database or the metadata store server.
disable_upgrade_migration: if set to True, the library does not upgrades
the db schema and migrates all data if it connects to an old version
backend.
"""
if isinstance(config, metadata_store_pb2.ConnectionConfig):
self._using_db_connection = True
migration_options = metadata_store_pb2.MigrationOptions()
migration_options.disable_upgrade_migration = disable_upgrade_migration
self._metadata_store = metadata_store_serialized.CreateMetadataStore(
config.SerializeToString(),
migration_options.SerializeToString())
# If you remove this line, errors are not thrown correctly.
logging.log(logging.INFO,
'MetadataStore with DB connection initialized')
return
if not isinstance(config,
metadata_store_pb2.MetadataStoreClientConfig):
raise ValueError('MetadataStore is expecting either'
'metadata_store_pb2.ConnectionConfig or'
'metadata_store_pb2.MetadataStoreClientConfig')
self._using_db_connection = False
target = ':'.join([config.host, str(config.port)])
channel = self._get_channel(config, target)
self._metadata_store_stub = (
metadata_store_service_pb2_grpc.MetadataStoreServiceStub(channel))
logging.log(logging.INFO,
'MetadataStore with gRPC connection initialized')
