def __init__(self, model=None, config=None):
self.model = model
self.structure = _Component("start", -1)
self.size = 0
if config:
self.config = NestedNamespace(parse_yaml(config))
else:
self.config = config
if self.model:
now = dt.datetime.now().strftime("%y%m%d_%h%m%s")
self.job_id = "{}_{}".format(self.model.model["name"], now)
def train(self, tune=False, blocking=True, wait_interval=60):
"""Trains on CAIP.
Args:
tune: train with hyperparameter tuning if true.
blocking: true if the function should exit only once the job
completes.
wait_interval: if blocking, how often the job state should be
checked.
Returns:
job_id: a CAIP job id.
"""
now = dt.datetime.now().strftime("%Y%m%d_%H%M%S")
job_id = "train_{}_{}".format(self.model["name"], now)
package_uri = self.upload_trainer_dist()
jobs_client = self.ml_client.projects().jobs()
body = {
"jobId": job_id,
"trainingInput": {
"scaleTier": self.scale_tier,
"packageUris": [package_uri],
"pythonModule": "trainer.task",
"args": [
"--model_dir",
self.get_model_dir(),
],
"jobDir": self.get_job_dir(),
"region": self.region,
"runtimeVersion": self.runtime_version,
"pythonVersion": self.python_version,
},
}
# TODO(smhosein): should we handle custom scale tiers?
if tune:
hp_config = parse_yaml(self.model_params["hyperparam_config"])
hyperparams = hp_config["trainingInput"]["hyperparameters"]
body["trainingInput"]["hyperparameters"] = hyperparams
request = jobs_client.create(parent=self.get_parent(), body=body)
self._call_ml_client(request)
if blocking:
self._wait_until_done(job_id, wait_interval)
return job_id
def generate_component(config, name, template_spec='./component_spec.yaml'):
"""Generate the component files from the templates."""
template_spec_path = path.join(path.dirname(__file__), template_spec)
output_spec = parse_yaml(template_spec_path)
current_spec = output_spec[name]
loader = jinja.PackageLoader('ml_pipeline_gen',
current_spec['template_dir'])
env = jinja.Environment(loader=loader,
trim_blocks=True,
lstrip_blocks='True')
template_file_list = current_spec['files']
for template in template_file_list:
template_in = env.get_template(template['input'])
template_out = template_in.render(config=config)
output_file = path.join(config.output_package, template['output'])
pathlib.Path(output_file).parent.mkdir(parents=True, exist_ok=True)
with open(output_file, 'w') as f:
f.write(template_out)
def __init__(self, config_path, framework):
config = parse_yaml(config_path)
self._set_config(config)
self.ml_client = discovery.build("ml", "v1")
self.framework = framework