def test_training_job():
rc = training_job.TrainingJobResourceConfig(
instance_type="test_type",
instance_count=10,
volume_size_in_gb=25,
distributed_protocol=training_job.DistributedProtocol.MPI,
)
alg = training_job.AlgorithmSpecification(
algorithm_name=training_job.AlgorithmName.CUSTOM,
algorithm_version="",
input_mode=training_job.InputMode.FILE,
input_content_type=training_job.InputContentType.TEXT_CSV,
)
tj = training_job.TrainingJob(
training_job_resource_config=rc,
algorithm_specification=alg,
)
tj2 = training_job.TrainingJob.from_flyte_idl(tj.to_flyte_idl())
# checking tj == tj2 would return false because we don't have the __eq__ magic method defined
assert tj.training_job_resource_config.instance_type == tj2.training_job_resource_config.instance_type
assert tj.training_job_resource_config.instance_count == tj2.training_job_resource_config.instance_count
assert tj.training_job_resource_config.distributed_protocol == tj2.training_job_resource_config.distributed_protocol
assert tj.training_job_resource_config.volume_size_in_gb == tj2.training_job_resource_config.volume_size_in_gb
assert tj.algorithm_specification.algorithm_name == tj2.algorithm_specification.algorithm_name
assert tj.algorithm_specification.algorithm_version == tj2.algorithm_specification.algorithm_version
assert tj.algorithm_specification.input_mode == tj2.algorithm_specification.input_mode
assert tj.algorithm_specification.input_content_type == tj2.algorithm_specification.input_content_type
def test_hyperparameter_tuning_job():
rc = training_job.TrainingJobResourceConfig(
instance_type="test_type",
instance_count=10,
volume_size_in_gb=25,
distributed_protocol=training_job.DistributedProtocol.MPI,
)
alg = training_job.AlgorithmSpecification(
algorithm_name=training_job.AlgorithmName.CUSTOM,
algorithm_version="",
input_mode=training_job.InputMode.FILE,
input_content_type=training_job.InputContentType.TEXT_CSV,
)
tj = training_job.TrainingJob(training_job_resource_config=rc, algorithm_specification=alg,)
hpo = hpo_job.HyperparameterTuningJob(max_number_of_training_jobs=10, max_parallel_training_jobs=5, training_job=tj)
hpo2 = hpo_job.HyperparameterTuningJob.from_flyte_idl(hpo.to_flyte_idl())
assert hpo.max_number_of_training_jobs == hpo2.max_number_of_training_jobs
assert hpo.max_parallel_training_jobs == hpo2.max_parallel_training_jobs
assert (
hpo2.training_job.training_job_resource_config.instance_type
== hpo.training_job.training_job_resource_config.instance_type
)
assert (
hpo2.training_job.training_job_resource_config.instance_count
== hpo.training_job.training_job_resource_config.instance_count
)
assert (
hpo2.training_job.training_job_resource_config.distributed_protocol
== hpo.training_job.training_job_resource_config.distributed_protocol
)
assert (
hpo2.training_job.training_job_resource_config.volume_size_in_gb
== hpo.training_job.training_job_resource_config.volume_size_in_gb
)
assert (
hpo2.training_job.algorithm_specification.algorithm_name
== hpo.training_job.algorithm_specification.algorithm_name
)
assert (
hpo2.training_job.algorithm_specification.algorithm_version
== hpo.training_job.algorithm_specification.algorithm_version
)
assert hpo2.training_job.algorithm_specification.input_mode == hpo.training_job.algorithm_specification.input_mode
assert (
hpo2.training_job.algorithm_specification.input_content_type
== hpo.training_job.algorithm_specification.input_content_type
)
def test_metric_definition():
md = training_job.MetricDefinition(name="test-metric", regex="[a-zA-Z]*")
md2 = training_job.MetricDefinition.from_flyte_idl(md.to_flyte_idl())
assert md == md2
assert md2.name == "test-metric"
assert md2.regex == "[a-zA-Z]*"
def test_algorithm_specification():
case = unittest.TestCase()
alg_spec = training_job.AlgorithmSpecification(
algorithm_name=training_job.AlgorithmName.CUSTOM,
algorithm_version="v100",
input_mode=training_job.InputMode.FILE,
metric_definitions=[training_job.MetricDefinition(name="a", regex="b")],
input_content_type=training_job.InputContentType.TEXT_CSV,
)
alg_spec2 = training_job.AlgorithmSpecification.from_flyte_idl(alg_spec.to_flyte_idl())
assert alg_spec2.algorithm_name == training_job.AlgorithmName.CUSTOM
assert alg_spec2.algorithm_version == "v100"
assert alg_spec2.input_mode == training_job.InputMode.FILE
case.assertCountEqual(alg_spec.metric_definitions, alg_spec2.metric_definitions)
assert alg_spec == alg_spec2
def test_training_job():
rc = training_job.TrainingJobResourceConfig(
from flytekit.models.sagemaker import training_job as training_job_models
import tensorflow as tf
from flytekit.common import utils as _common_utils
@inputs(dummy_train_dataset=Types.Blob,
dummy_validation_dataset=Types.Blob,
my_input=Types.String)
@outputs(out_model=Types.Blob,
out=Types.Integer,
out_extra_output_file=Types.Blob)
@custom_training_job_task(
algorithm_specification=training_job_models.AlgorithmSpecification(
input_mode=training_job_models.InputMode.FILE,
algorithm_name=training_job_models.AlgorithmName.CUSTOM,
algorithm_version="",
input_content_type=training_job_models.InputContentType.TEXT_CSV,
),
training_job_resource_config=training_job_models.TrainingJobResourceConfig(
instance_type="ml.m4.xlarge",
instance_count=1,
volume_size_in_gb=25,
))
def custom_training_task(wf_params, dummy_train_dataset,
dummy_validation_dataset, my_input, out_model, out,
out_extra_output_file):
with _common_utils.AutoDeletingTempDir("output_dir") as output_dir:
wf_params.logging.info("My printed value: {}".format(my_input))
wf_params.logging.info(
"My dummy train_dataset: {}".format(dummy_train_dataset))
wf_params.logging.info(