def test_start_training__output_dir(output_dir: Optional[str],
expected_output_dir: str,
mocker: MockFixture):
"""
Tests output_dir parameter. If it is provided, the result model directory is created under it.
Otherwise, it is created in a default output path from config
"""
trainer = K8sTrainer(model_training=ModelTraining(spec=ModelTrainingSpec(
model=ModelIdentity())))
mocker.patch.object(training, 'create_mt_config_file')
config_mock = mocker.patch.object(training, 'config')
os_makedirs_mock = mocker.patch.object(training.os, 'makedirs')
compile_artifact_name_template_mock = mocker.patch.object(
training, 'compile_artifact_name_template')
launch_training_container_mock = mocker.patch.object(
training, 'launch_training_container')
launch_gppi_validation_container_mock = mocker.patch.object(
training, 'launch_gppi_validation_container')
config_mock.LOCAL_MODEL_OUTPUT_DIR = DEFAULT_OUTPUT_DIR
compile_artifact_name_template_mock.return_value = 'model_dir_name'
start_train(trainer, output_dir)
expected_model_dir_path = os.path.join(expected_output_dir,
'model_dir_name')
os_makedirs_mock.assert_called_with(expected_model_dir_path, exist_ok=ANY)
launch_training_container_mock.assert_called_with(trainer,
expected_model_dir_path)
launch_gppi_validation_container_mock.assert_called_with(
trainer, expected_model_dir_path)
def test_local_training_relative_output_dir(mocker: MockFixture, cli_runner: CliRunner):
"""
Tests issue #208 - Converting relative path to output-dir to absolute,
because Docker requires mount paths to be absolute
:param mocker: mocker fixture
:param cli_runner: Click runner fixture
"""
trainer_mock = mocker.patch.object(training, 'K8sTrainer', autospec=True).return_value
trainer_mock.model_training.spec.model.artifact_name_template = 'model_dir_template'
api_client = Mock()
mocker.patch.object(training_sdk, 'create_mt_config_file')
mocker.patch.object(training_sdk, 'stream_container_logs')
mocker.patch.object(training_sdk, 'raise_error_if_container_failed')
docker_mock: Mock = mocker.patch.object(training_sdk.docker, 'from_env').return_value
docker_mock.api.inspect_container = Mock(return_value={})
with tempfile.TemporaryDirectory(dir=os.curdir) as temp_dir:
temp_dir_path = pathlib.Path(temp_dir)
training_yml = temp_dir_path / 'training.yml'
training_yml.write_text(json.dumps(
{**ModelTraining(id='training1', spec=ModelTrainingSpec(toolchain='toolchain1')).to_dict(),
**{'kind': 'ModelTraining'}}))
toolchain_yml = temp_dir_path / 'toolchain.yml'
toolchain_yml.write_text(json.dumps({**ToolchainIntegration(id='toolchain1').to_dict(),
**{'kind': 'ToolchainIntegration'}}))
temp_dir_relative_path: str = os.path.relpath(temp_dir_path)
result: Result = cli_runner.invoke(
training.training_group,
['run', '--output-dir', temp_dir_relative_path, '--manifest-file', str(training_yml),
'--manifest-file', str(toolchain_yml), '--id', 'training1'],
obj=api_client)
assert result.exit_code == 0, f'command invocation ended with exit code {result.exit_code}'
assert result.exception is None, f'command invocation ended with exception {result.exception}'
abs_path = os.path.abspath(temp_dir)
for call in docker_mock.containers.run.call_args_list:
call_kwargs = call[1]
mounts: List[Mount] = call_kwargs.get('mounts', [])
for mount in filter(lambda m: m.get('Target') == MODEL_OUTPUT_CONTAINER_PATH, mounts):
assert os.path.dirname(mount.get('Source')) == abs_path
def test_start_training__artifact_name_template(template: Optional[str],
expected_template: str,
mocker: MockFixture):
"""
Tests artifactNameTemplate property from training configuration. If it is provided,
the result model directory is named according to it. Otherwise, the default template is used.
If artifact name template ends with .zip it is trimmed.
"""
model = ModelIdentity(artifact_name_template=template)
trainer = K8sTrainer(model_training=ModelTraining(spec=ModelTrainingSpec(
model=model)))
mocker.patch.object(training, 'create_mt_config_file')
os_makedirs_mock = mocker.patch.object(training.os, 'makedirs')
compile_artifact_name_template_mock = mocker.patch.object(
training, 'compile_artifact_name_template')
launch_training_container_mock = mocker.patch.object(
training, 'launch_training_container')
launch_gppi_validation_container_mock = mocker.patch.object(
training, 'launch_gppi_validation_container')
compile_artifact_name_template_mock.return_value = 'compiled_model_dir_name'
output_dir = '/test/output'
start_train(trainer, output_dir)
compile_artifact_name_template_mock.assert_called_with(
go_template=expected_template, context=ANY)
expected_model_dir_path = os.path.join(output_dir,
'compiled_model_dir_name')
os_makedirs_mock.assert_called_with(expected_model_dir_path, exist_ok=ANY)
launch_training_container_mock.assert_called_with(trainer,
expected_model_dir_path)
launch_gppi_validation_container_mock.assert_called_with(
trainer, expected_model_dir_path)
type="gcs",
uri=f'gs://{wine_bucket}/data/wine-quality.csv',
region=gcp_project,
))
training_id = "airlfow-wine"
training = ModelTraining(
id=training_id,
spec=ModelTrainingSpec(
model=ModelIdentity(name="wine", version="1.0"),
toolchain="mlflow",
entrypoint="main",
work_dir="mlflow/sklearn/wine",
hyper_parameters={"alpha": "1.0"},
data=[
DataBindingDir(conn_name='wine',
local_path='mlflow/sklearn/wine/wine-quality.csv'),
],
resources=ResourceRequirements(requests=ResourceList(cpu="2024m",
memory="2024Mi"),
limits=ResourceList(cpu="2024m",
memory="2024Mi")),
vcs_name="odahu-flow-examples"),
)
packaging_id = "airlfow-wine"
packaging = ModelPackaging(
id=packaging_id,
spec=ModelPackagingSpec(
targets=[Target(name="docker-push", connection_name="docker-ci")],
integration_name="docker-rest"),
'email_on_retry': False,
'end_date': datetime(2099, 12, 31)
}
api_connection_id = "odahuflow_api"
model_connection_id = "odahuflow_model"
training_id = "airflow-tensorflow"
training = ModelTraining(
id=training_id,
spec=ModelTrainingSpec(model=ModelIdentity(name="tensorflow",
version="1.0"),
toolchain="mlflow",
entrypoint="main",
work_dir="mlflow/tensorflow/example",
resources=ResourceRequirements(
requests=ResourceList(cpu="2024m",
memory="2024Mi"),
limits=ResourceList(cpu="2024m",
memory="2024Mi")),
vcs_name="odahu-flow-examples"),
)
packaging_id = "airflow-tensorflow"
packaging = ModelPackaging(
id=packaging_id,
spec=ModelPackagingSpec(
targets=[Target(name="docker-push", connection_name="docker-ci")],
integration_name="docker-rest"),
)
id=ENTITY_ID,
spec=ConnectionSpec(
key_secret="mock-key-secret",
uri="mock-url",
),
),
connection.connection,
'Connection',
)
TRAINING = EntityTestData(
ModelTrainingClient(),
ModelTraining(
id=ENTITY_ID,
spec=ModelTrainingSpec(work_dir="/1/2/3",
model=ModelIdentity(name="name",
version="version")),
status=ModelTrainingStatus(state=SUCCEEDED_STATE)), training.training,
'ModelTraining')
TOOLCHAIN = EntityTestData(
ToolchainIntegrationClient(),
ToolchainIntegration(
id=ENTITY_ID,
spec=ToolchainIntegrationSpec(
default_image="mock-image",
entrypoint="default-entrypoint",
),
),
toolchain_integration.toolchain_integration,
'ToolchainIntegration',