def __exit__(self, exc_type, exc_value, exc_traceback):
self._end_time = datetime.datetime.now(dateutil.tz.tzlocal())
if not self._in_sagemaker_job:
self.trial_component.end_time = self._end_time
if exc_value:
self.trial_component.status = api_types.TrialComponentStatus(
primary_status="Failed", message=str(exc_value))
else:
self.trial_component.status = api_types.TrialComponentStatus(
primary_status="Completed")
self.close()
def __enter__(self):
self._start_time = datetime.datetime.now(dateutil.tz.tzlocal())
if not self._in_sagemaker_job:
self.trial_component.start_time = self._start_time
self.trial_component.status = api_types.TrialComponentStatus(
primary_status="InProgress")
return self
def test_save(trial_component_obj, sagemaker_boto_client):
trial_component_obj.display_name = str(uuid.uuid4())
trial_component_obj.status = api_types.TrialComponentStatus(
primary_status="InProgress", message="Message")
trial_component_obj.start_time = datetime.datetime.now(
datetime.timezone.utc) - datetime.timedelta(days=1)
trial_component_obj.end_time = datetime.datetime.now(datetime.timezone.utc)
trial_component_obj.parameters = {"foo": "bar", "whizz": 100.1}
trial_component_obj.input_artifacts = {
"snizz":
api_types.TrialComponentArtifact(value="s3:/foo/bar",
media_type="text/plain")
}
trial_component_obj.output_artifacts = {
"fly":
api_types.TrialComponentArtifact(value="s3:/sky/far",
media_type="away/tomorrow")
}
trial_component_obj.save()
loaded = trial_component.TrialComponent.load(
trial_component_name=trial_component_obj.trial_component_name,
sagemaker_boto_client=sagemaker_boto_client)
assert trial_component_obj.trial_component_name == loaded.trial_component_name
assert trial_component_obj.status == loaded.status
assert trial_component_obj.start_time - loaded.start_time < datetime.timedelta(
seconds=1)
assert trial_component_obj.end_time - loaded.end_time < datetime.timedelta(
seconds=1)
assert trial_component_obj.parameters == loaded.parameters
assert trial_component_obj.input_artifacts == loaded.input_artifacts
assert trial_component_obj.output_artifacts == loaded.output_artifacts
def test_save(trial_component_obj, sagemaker_boto_client):
trial_component_obj.display_name = str(uuid.uuid4())
trial_component_obj.status = api_types.TrialComponentStatus(primary_status='InProgress', message='Message')
trial_component_obj.start_time = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(days=1)
trial_component_obj.end_time = datetime.datetime.now(datetime.timezone.utc)
trial_component_obj.parameters = {
'foo': 'bar',
'whizz': 100.1
}
trial_component_obj.input_artifacts = {
'snizz': api_types.TrialComponentArtifact(value='s3:/foo/bar', media_type='text/plain')
}
trial_component_obj.output_artifacts = {
'fly': api_types.TrialComponentArtifact(value='s3:/sky/far', media_type='away/tomorrow')
}
trial_component_obj.save()
loaded = trial_component.TrialComponent.load(trial_component_name = trial_component_obj.trial_component_name,
sagemaker_boto_client=sagemaker_boto_client)
assert trial_component_obj.trial_component_name == loaded.trial_component_name
assert trial_component_obj.status == loaded.status
assert trial_component_obj.start_time - loaded.start_time < datetime.timedelta(seconds=1)
assert trial_component_obj.end_time - loaded.end_time < datetime.timedelta(seconds=1)
assert trial_component_obj.parameters == loaded.parameters
assert trial_component_obj.input_artifacts == loaded.input_artifacts
assert trial_component_obj.output_artifacts == loaded.output_artifacts
def test_load(sagemaker_boto_client):
now = datetime.datetime.now(datetime.timezone.utc)
sagemaker_boto_client.describe_trial_component.return_value = {
"TrialComponentArn": "A",
"TrialComponentName": "B",
"DisplayName": "C",
"Status": {"PrimaryStatus": "InProgress", "Message": "D"},
"Parameters": {"E": {"NumberValue": 1.0}, "F": {"StringValue": "G"}},
"InputArtifacts": {"H": {"Value": "s3://foo/bar", "MediaType": "text/plain"}},
"OutputArtifacts": {"I": {"Value": "s3://whizz/bang", "MediaType": "text/plain"}},
"Metrics": [
{
"MetricName": "J",
"Count": 1,
"Min": 1.0,
"Max": 2.0,
"Avg": 3.0,
"StdDev": 4.0,
"SourceArn": "K",
"Timestamp": now,
}
],
}
obj = trial_component.TrialComponent.load(trial_component_name="foo", sagemaker_boto_client=sagemaker_boto_client)
sagemaker_boto_client.describe_trial_component.assert_called_with(TrialComponentName="foo")
assert "A" == obj.trial_component_arn
assert "B" == obj.trial_component_name
assert "C" == obj.display_name
assert api_types.TrialComponentStatus(primary_status="InProgress", message="D") == obj.status
assert {"E": 1.0, "F": "G"} == obj.parameters
assert {"H": api_types.TrialComponentArtifact(value="s3://foo/bar", media_type="text/plain")}
assert {"I": api_types.TrialComponentArtifact(value="s3://whizz/bang", media_type="text/plain")}
assert [
api_types.TrialComponentMetricSummary(
metric_name="J", count=1, min=1.0, max=2.0, avg=3.0, std_dev=4.0, source_arn="K", timestamp=now
)
]
def test_list(sagemaker_boto_client):
start_time = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(hours=1)
end_time = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(hours=2)
creation_time = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(hours=3)
last_modified_time = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(hours=4)
sagemaker_boto_client.list_trial_components.side_effect = [
{
"TrialComponentSummaries": [
{
"TrialComponentName": "A" + str(i),
"TrialComponentArn": "B" + str(i),
"DisplayName": "C" + str(i),
"SourceArn": "D" + str(i),
"Status": {"PrimaryStatus": "InProgress", "Message": "E" + str(i)},
"StartTime": start_time + datetime.timedelta(hours=i),
"EndTime": end_time + datetime.timedelta(hours=i),
"CreationTime": creation_time + datetime.timedelta(hours=i),
"LastModifiedTime": last_modified_time + datetime.timedelta(hours=i),
"LastModifiedBy": {},
}
for i in range(10)
],
"NextToken": "100",
},
{
"TrialComponentSummaries": [
{
"TrialComponentName": "A" + str(i),
"TrialComponentArn": "B" + str(i),
"DisplayName": "C" + str(i),
"SourceArn": "D" + str(i),
"Status": {"PrimaryStatus": "InProgress", "Message": "E" + str(i)},
"StartTime": start_time + datetime.timedelta(hours=i),
"EndTime": end_time + datetime.timedelta(hours=i),
"CreationTime": creation_time + datetime.timedelta(hours=i),
"LastModifiedTime": last_modified_time + datetime.timedelta(hours=i),
"LastModifiedBy": {},
}
for i in range(10, 20)
]
},
]
expected = [
api_types.TrialComponentSummary(
trial_component_name="A" + str(i),
trial_component_arn="B" + str(i),
display_name="C" + str(i),
source_arn="D" + str(i),
status=api_types.TrialComponentStatus(primary_status="InProgress", message="E" + str(i)),
start_time=start_time + datetime.timedelta(hours=i),
end_time=end_time + datetime.timedelta(hours=i),
creation_time=creation_time + datetime.timedelta(hours=i),
last_modified_time=last_modified_time + datetime.timedelta(hours=i),
last_modified_by={},
)
for i in range(20)
]
result = list(
trial_component.TrialComponent.list(
sagemaker_boto_client=sagemaker_boto_client,
source_arn="foo",
sort_by="CreationTime",
sort_order="Ascending",
)
)
assert expected == result
expected_calls = [
unittest.mock.call(SortBy="CreationTime", SortOrder="Ascending", SourceArn="foo"),
unittest.mock.call(NextToken="100", SortBy="CreationTime", SortOrder="Ascending", SourceArn="foo"),
]
assert expected_calls == sagemaker_boto_client.list_trial_components.mock_calls
def test_load(sagemaker_boto_client):
now = datetime.datetime.now(datetime.timezone.utc)
sagemaker_boto_client.describe_trial_component.return_value = {
'TrialComponentArn':
'A',
'TrialComponentName':
'B',
'DisplayName':
'C',
'Status': {
'PrimaryStatus': 'InProgress',
'Message': 'D'
},
'Parameters': {
'E': {
'NumberValue': 1.0
},
'F': {
'StringValue': 'G'
}
},
'InputArtifacts': {
'H': {
'Value': 's3://foo/bar',
'MediaType': 'text/plain'
}
},
'OutputArtifacts': {
'I': {
'Value': 's3://whizz/bang',
'MediaType': 'text/plain'
}
},
'Metrics': [{
'MetricName': 'J',
'Count': 1,
'Min': 1.0,
'Max': 2.0,
'Avg': 3.0,
'StdDev': 4.0,
'SourceArn': 'K',
'Timestamp': now
}]
}
obj = trial_component.TrialComponent.load(
trial_component_name='foo',
sagemaker_boto_client=sagemaker_boto_client)
sagemaker_boto_client.describe_trial_component.assert_called_with(
TrialComponentName='foo')
assert 'A' == obj.trial_component_arn
assert 'B' == obj.trial_component_name
assert 'C' == obj.display_name
assert api_types.TrialComponentStatus(primary_status='InProgress',
message='D') == obj.status
assert {'E': 1.0, 'F': 'G'} == obj.parameters
assert {
'H':
api_types.TrialComponentArtifact(value='s3://foo/bar',
media_type='text/plain')
}
assert {
'I':
api_types.TrialComponentArtifact(value='s3://whizz/bang',
media_type='text/plain')
}
assert [
api_types.TrialComponentMetricSummary(metric_name='J',
count=1,
min=1.0,
max=2.0,
avg=3.0,
std_dev=4.0,
source_arn='K',
timestamp=now)
]
def test_list(sagemaker_boto_client):
start_time = datetime.datetime.now(
datetime.timezone.utc) + datetime.timedelta(hours=1)
end_time = datetime.datetime.now(
datetime.timezone.utc) + datetime.timedelta(hours=2)
creation_time = datetime.datetime.now(
datetime.timezone.utc) + datetime.timedelta(hours=3)
last_modified_time = datetime.datetime.now(
datetime.timezone.utc) + datetime.timedelta(hours=4)
sagemaker_boto_client.list_trial_components.side_effect = [
{
'TrialComponentSummaries': [{
'TrialComponentName':
'A' + str(i),
'TrialComponentArn':
'B' + str(i),
'DisplayName':
'C' + str(i),
'SourceArn':
'D' + str(i),
'Status': {
'PrimaryStatus': 'InProgress',
'Message': 'E' + str(i)
},
'StartTime':
start_time + datetime.timedelta(hours=i),
'EndTime':
end_time + datetime.timedelta(hours=i),
'CreationTime':
creation_time + datetime.timedelta(hours=i),
'LastModifiedTime':
last_modified_time + datetime.timedelta(hours=i),
'LastModifiedBy': {}
} for i in range(10)],
'NextToken':
'100'
},
{
'TrialComponentSummaries': [{
'TrialComponentName':
'A' + str(i),
'TrialComponentArn':
'B' + str(i),
'DisplayName':
'C' + str(i),
'SourceArn':
'D' + str(i),
'Status': {
'PrimaryStatus': 'InProgress',
'Message': 'E' + str(i)
},
'StartTime':
start_time + datetime.timedelta(hours=i),
'EndTime':
end_time + datetime.timedelta(hours=i),
'CreationTime':
creation_time + datetime.timedelta(hours=i),
'LastModifiedTime':
last_modified_time + datetime.timedelta(hours=i),
'LastModifiedBy': {}
} for i in range(10, 20)]
},
]
expected = [
api_types.TrialComponentSummary(
trial_component_name='A' + str(i),
trial_component_arn='B' + str(i),
display_name='C' + str(i),
source_arn='D' + str(i),
status=api_types.TrialComponentStatus(primary_status='InProgress',
message='E' + str(i)),
start_time=start_time + datetime.timedelta(hours=i),
end_time=end_time + datetime.timedelta(hours=i),
creation_time=creation_time + datetime.timedelta(hours=i),
last_modified_time=last_modified_time +
datetime.timedelta(hours=i),
last_modified_by={}) for i in range(20)
]
result = list(
trial_component.TrialComponent.list(
sagemaker_boto_client=sagemaker_boto_client,
source_arn='foo',
sort_by='CreationTime',
sort_order='Ascending'))
assert expected == result
expected_calls = [
unittest.mock.call(SortBy='CreationTime',
SortOrder='Ascending',
SourceArn='foo'),
unittest.mock.call(NextToken='100',
SortBy='CreationTime',
SortOrder='Ascending',
SourceArn='foo')
]
assert expected_calls == sagemaker_boto_client.list_trial_components.mock_calls