119 lines
3.9 KiB
Python
119 lines
3.9 KiB
Python
# Copyright 2018-2022 The Kubeflow Authors
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from itertools import chain
|
|
import os
|
|
|
|
import click
|
|
import kfp
|
|
from kfp import client
|
|
from kfp.cli import component
|
|
from kfp.cli import diagnose_me_cli
|
|
from kfp.cli import dsl
|
|
from kfp.cli import experiment
|
|
from kfp.cli import pipeline
|
|
from kfp.cli import recurring_run
|
|
from kfp.cli import run
|
|
from kfp.cli.output import OutputFormat
|
|
from kfp.cli.utils import aliased_plurals_group
|
|
from kfp.cli.utils import parsing
|
|
|
|
COMMANDS = {
|
|
'client': {
|
|
run.run, recurring_run.recurring_run, experiment.experiment,
|
|
pipeline.pipeline
|
|
},
|
|
'no_client': {diagnose_me_cli.diagnose_me, component.component, dsl.dsl}
|
|
}
|
|
|
|
PROGRAM_NAME = 'kfp'
|
|
|
|
SHELL_FILES = {
|
|
'bash': ['.bashrc'],
|
|
'zsh': ['.zshrc'],
|
|
'fish': ['.config', 'fish', 'completions', f'{PROGRAM_NAME}.fish']
|
|
}
|
|
|
|
|
|
def _create_completion(shell: str) -> str:
|
|
return f'eval "$(_{PROGRAM_NAME.upper()}_COMPLETE={shell}_source {PROGRAM_NAME})"'
|
|
|
|
|
|
def _install_completion(shell: str) -> None:
|
|
completion_statement = _create_completion(shell)
|
|
source_file = os.path.join(os.path.expanduser('~'), *SHELL_FILES[shell])
|
|
with open(source_file, 'a') as f:
|
|
f.write('\n' + completion_statement + '\n')
|
|
|
|
|
|
@click.group(
|
|
name=PROGRAM_NAME,
|
|
cls=aliased_plurals_group.AliasedPluralsGroup, # type: ignore
|
|
commands=list(chain.from_iterable(COMMANDS.values())), # type: ignore
|
|
invoke_without_command=True)
|
|
@click.option(
|
|
'--show-completion',
|
|
type=click.Choice(list(SHELL_FILES.keys())),
|
|
default=None)
|
|
@click.option(
|
|
'--install-completion',
|
|
type=click.Choice(list(SHELL_FILES.keys())),
|
|
default=None)
|
|
@click.option('--endpoint', help=parsing.get_param_descr(client.Client, 'host'))
|
|
@click.option(
|
|
'--iap-client-id', help=parsing.get_param_descr(client.Client, 'client_id'))
|
|
@click.option(
|
|
'-n',
|
|
'--namespace',
|
|
default='kubeflow',
|
|
show_default=True,
|
|
help=parsing.get_param_descr(client.Client, 'namespace'))
|
|
@click.option(
|
|
'--other-client-id',
|
|
help=parsing.get_param_descr(client.Client, 'other_client_id'))
|
|
@click.option(
|
|
'--other-client-secret',
|
|
help=parsing.get_param_descr(client.Client, 'other_client_secret'))
|
|
@click.option(
|
|
'--output',
|
|
type=click.Choice(list(map(lambda x: x.name, OutputFormat))),
|
|
default=OutputFormat.table.name,
|
|
show_default=True,
|
|
help='The formatting style for command output.')
|
|
@click.pass_context
|
|
@click.version_option(version=kfp.__version__, message='%(prog)s %(version)s')
|
|
def cli(ctx: click.Context, endpoint: str, iap_client_id: str, namespace: str,
|
|
other_client_id: str, other_client_secret: str, output: OutputFormat,
|
|
show_completion: str, install_completion: str):
|
|
"""Kubeflow Pipelines CLI."""
|
|
if show_completion:
|
|
click.echo(_create_completion(show_completion))
|
|
return
|
|
if install_completion:
|
|
_install_completion(install_completion)
|
|
return
|
|
|
|
client_commands = set(
|
|
chain.from_iterable([
|
|
(command.name, f'{command.name}s')
|
|
for command in COMMANDS['client'] # type: ignore
|
|
]))
|
|
if ctx.invoked_subcommand not in client_commands:
|
|
# Do not create a client for these subcommands
|
|
return
|
|
ctx.obj['client'] = client.Client(endpoint, iap_client_id, namespace,
|
|
other_client_id, other_client_secret)
|
|
ctx.obj['namespace'] = namespace
|
|
ctx.obj['output'] = output
|