def test_57_put_files_uuid_replace_locations(self):
"""Test that a file can replace with the same location."""
self.start_server()
token = self.get_token()
r = RestClient(self.address, token, timeout=1, retries=1)
# define the files to be created
metadata = {
'logical_name': '/blah/data/exp/IceCube/blah.dat',
'checksum': {'sha512': hex('foo bar')},
'file_size': 1,
u'locations': [{u'site': u'WIPAC', u'path': u'/blah/data/exp/IceCube/blah.dat'}]
}
metadata2 = {
'logical_name': '/blah/data/exp/IceCube/blah.dat',
'checksum': {'sha512': hex('foo bar2')},
'file_size': 2,
u'locations': [{u'site': u'WIPAC', u'path': u'/blah/data/exp/IceCube/blah.dat'}]
}
# create the first file; should be OK
data = r.request_seq('POST', '/api/files', metadata)
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('file', data)
url = data['file']
uid = url.split('/')[-1]
# try to replace the first file with the second; should be OK
data = r.request_seq('PUT', '/api/files/'+uid, metadata2)
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('logical_name', data)
async def _do_work_claim(self) -> bool:
"""Claim a bundle and perform work on it."""
# 1. Ask the LTA DB for the next Bundle to be built
# configure a RestClient to talk to the File Catalog
fc_rc = RestClient(self.file_catalog_rest_url,
token=self.file_catalog_rest_token,
timeout=self.work_timeout_seconds,
retries=self.work_retries)
# configure a RestClient to talk to the LTA DB
lta_rc = RestClient(self.lta_rest_url,
token=self.lta_rest_token,
timeout=self.work_timeout_seconds,
retries=self.work_retries)
self.logger.info("Asking the LTA DB for a Bundle to build.")
pop_body = {"claimant": f"{self.name}-{self.instance_uuid}"}
response = await lta_rc.request(
'POST',
f'/Bundles/actions/pop?source={self.source_site}&dest={self.dest_site}&status={self.input_status}',
pop_body)
self.logger.info(f"LTA DB responded with: {response}")
bundle = response["bundle"]
if not bundle:
self.logger.info(
"LTA DB did not provide a Bundle to build. Going on vacation.")
return False
# process the Bundle that we were given
try:
await self._do_work_bundle(fc_rc, lta_rc, bundle)
except Exception as e:
await self._quarantine_bundle(lta_rc, bundle, f"{e}")
raise e
# signal the work was processed successfully
return True
def test_54_patch_files_uuid_replace_logical_name(self):
"""Test that a file can be updated with the same logical_name."""
self.start_server()
token = self.get_token()
r = RestClient(self.address, token, timeout=1, retries=1)
# define the file to be created
metadata = {
'logical_name': '/blah/data/exp/IceCube/blah.dat',
'checksum': {'sha512': hex('foo bar')},
'file_size': 1,
u'locations': [{u'site': u'WIPAC', u'path': u'/blah/data/exp/IceCube/blah.dat'}]
}
# this is a PATCH to metadata; matches the old logical_name
patch1 = {
'logical_name': '/blah/data/exp/IceCube/blah.dat',
'checksum': {'sha512': hex('foo bar2')},
'file_size': 2,
u'locations': [{u'site': u'WIPAC', u'path': u'/blah/data/exp/IceCube/blah.dat'}]
}
# create the file; should be OK
data = r.request_seq('POST', '/api/files', metadata)
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('file', data)
url = data['file']
uid = url.split('/')[-1]
# try to update the file with a patch; should be OK
data = r.request_seq('PATCH', '/api/files/' + uid, patch1)
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('logical_name', data)
def test_15_files_auth(self):
appkey = 'secret2'
self.edit_config({
'auth':{
'secret': 'secret',
'expiration': 82400,
}
})
self.start_server()
token = self.get_token()
r = RestClient(self.address, token, timeout=1, retries=1)
metadata = {
'logical_name': 'blah',
'checksum': {'sha512':hex('foo bar')},
'file_size': 1,
u'locations': [{u'site':u'test',u'path':u'blah.dat'}]
}
r2 = RestClient(self.address, 'blah', timeout=1, retries=1)
with self.assertRaises(Exception):
r2.request_seq('POST', '/api/files', metadata)
data = r.request_seq('POST', '/api/files', metadata)
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('file', data)
url = data['file']
def test_71_snapshot_find(self):
self.start_server()
token = self.get_token()
r = RestClient(self.address, token, timeout=1, retries=1)
metadata = {
'collection_name': 'blah',
'owner': 'foo',
}
data = r.request_seq('POST', '/api/collections', metadata)
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('collection', data)
url = data['collection']
uid = url.split('/')[-1]
data = r.request_seq('GET', '/api/collections/' + uid)
data = r.request_seq('POST', '/api/collections/{}/snapshots'.format(uid))
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('snapshot', data)
url = data['snapshot']
snap_uid = url.split('/')[-1]
data = r.request_seq('GET', '/api/snapshots/' + snap_uid)
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('files',data)
self.assertEqual(data['files'], [])
def test_10_files(self):
self.start_server()
token = self.get_token()
r = RestClient(self.address, token, timeout=1, retries=1)
metadata = {
'logical_name': 'blah',
'checksum': {'sha512':hex('foo bar')},
'file_size': 1,
u'locations': [{u'site':u'test',u'path':u'blah.dat'}]
}
data = r.request_seq('POST', '/api/files', metadata)
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('file', data)
url = data['file']
uid = url.split('/')[-1]
data = r.request_seq('GET', '/api/files')
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('files', data)
self.assertEqual(len(data['files']), 1)
self.assertTrue(any(uid == f['uuid'] for f in data['files']))
for m in ('PUT','DELETE','PATCH'):
with self.assertRaises(Exception):
r.request_seq(m, '/api/files')
def test_101_request_stream() -> None:
"""Test `request_stream()` when there's no response."""
mock_url = "http://test"
rpc = RestClient(mock_url, "passkey", timeout=1)
empty_streams = [
[b"\n"],
[],
[b"\n", b"\r\n", b"\n"],
[b" \n"],
[b" "],
[b"\t"],
]
for expected_stream in empty_streams:
# test multiple times
for test_num in range(2):
print(f"\niteration #{test_num}")
HTTPretty.register_uri(
HTTPretty.POST,
mock_url + "/stream/no-resp/",
body=expected_stream,
streaming=True,
)
response_stream = rpc.request_stream("POST", "/stream/no-resp/",
{})
never_entered = True
with _in_time(10, "Iterating by line is taking forever!"):
for _ in response_stream:
never_entered = False
assert never_entered
# now w/ chunk sizes
for chunk_size in [
None, -1, 0, 1, 2, 3, 4, 8, 9, 20, 100, 1024, 32768
]:
print(f"\nchunk_size: {chunk_size}")
HTTPretty.register_uri(
HTTPretty.POST,
mock_url + "/stream/no-resp/w/chunks",
body=expected_stream,
streaming=True,
)
response_stream = rpc.request_stream("POST",
"/stream/no-resp/w/chunks",
{},
chunk_size=chunk_size)
never_entered_w_chunks = True
with _in_time(10,
"Iterating by line is taking forever w/ chunks!"):
for _ in response_stream:
never_entered_w_chunks = False
assert never_entered_w_chunks
async def test_92_request_seq(requests_mock: Mock) -> None:
"""Test `request_seq()`."""
rpc = RestClient("http://test", "passkey", timeout=0.1)
def response(req: PreparedRequest, ctx: object) -> bytes: # pylint: disable=W0613
raise Exception()
requests_mock.post("/test", content=response)
with pytest.raises(Exception):
rpc.request_seq("POST", "test", {})
def test_80_snapshot_files(self):
self.start_server()
token = self.get_token()
r = RestClient(self.address, token, timeout=1, retries=1)
metadata = {
'collection_name': 'blah',
'owner': 'foo',
}
data = r.request_seq('POST', '/api/collections', metadata)
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('collection', data)
url = data['collection']
uid = url.split('/')[-1]
data = r.request_seq('GET', '/api/collections/' + uid)
data = r.request_seq('POST', '/api/collections/{}/snapshots'.format(uid))
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('snapshot', data)
url = data['snapshot']
snap_uid = url.split('/')[-1]
data = r.request_seq('GET', '/api/snapshots/{}/files'.format(snap_uid))
self.assertEqual(data['files'], [])
# add a file
metadata = {
'logical_name': 'blah',
'checksum': {'sha512':hex('foo bar')},
'file_size': 1,
u'locations': [{u'site':u'test',u'path':u'blah.dat'}]
}
data = r.request_seq('POST', '/api/files', metadata)
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('file', data)
url = data['file']
file_uid = url.split('/')[-1]
# old snapshot stays empty
data = r.request_seq('GET', '/api/snapshots/{}/files'.format(snap_uid))
self.assertEqual(data['files'], [])
# new snapshot should have file
data = r.request_seq('POST', '/api/collections/{}/snapshots'.format(uid))
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('snapshot', data)
url = data['snapshot']
snap_uid = url.split('/')[-1]
data = r.request_seq('GET', '/api/snapshots/{}/files'.format(snap_uid),
{'keys':'uuid|logical_name|checksum|locations'})
self.assertEqual(len(data['files']), 1)
self.assertEqual(data['files'][0]['uuid'], file_uid)
self.assertEqual(data['files'][0]['checksum'], metadata['checksum'])
def test_66_patch_files_uuid_locations_1xN(self):
"""Test locations uniqueness under 1xN multiplicity."""
self.start_server()
token = self.get_token()
r = RestClient(self.address, token, timeout=1, retries=1)
# define the locations to be tested
loc1a = {'site': 'WIPAC', 'path': '/data/test/exp/IceCube/foo.dat'}
loc1b = {'site': 'DESY', 'path': '/data/test/exp/IceCube/foo.dat'}
loc1c = {'site': 'NERSC', 'path': '/data/test/exp/IceCube/foo.dat'}
loc1d = {'site': 'OSG', 'path': '/data/test/exp/IceCube/foo.dat'}
locs3a = [loc1a, loc1b, loc1c]
locs3b = [loc1b, loc1c, loc1d]
locs3c = [loc1a, loc1b, loc1d]
# define the files to be created
metadata = {
'logical_name': '/blah/data/exp/IceCube/blah.dat',
'checksum': {'sha512': hex('foo bar')},
'file_size': 1,
u'locations': [loc1a]
}
metadata2 = {
'logical_name': '/blah/data/exp/IceCube/blah2.dat',
'checksum': {'sha512': hex('foo bar')},
'file_size': 1,
u'locations': [loc1b]
}
# this is a PATCH to metadata; steps on metadata2's location
patch1 = {
'logical_name': '/blah/data/exp/IceCube/blah.dat',
'checksum': {'sha512': hex('foo bar2')},
'file_size': 2,
u'locations': locs3c
}
# create the first file; should be OK
data = r.request_seq('POST', '/api/files', metadata)
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('file', data)
url = data['file']
uid = url.split('/')[-1]
# create the second file; should be OK
data = r.request_seq('POST', '/api/files', metadata2)
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('file', data)
# try to update the first file with a patch; should NOT be OK
with self.assertRaises(Exception):
r.request_seq('PATCH', '/api/files/' + uid, patch1)
def test_post_files(db_rc: RestClient) -> None:
"""Failure-test role authorization."""
post_body = {
"database": "test_histograms",
"collection": "collection_name",
"files": ["test.txt"],
}
with pytest.raises(requests.exceptions.HTTPError) as e:
db_rc.request_seq("POST", "/files/names", post_body)
assert e.response.status_code == 403 # Forbidden Error
db_rc.close()
def test_post_histo(db_rc: RestClient) -> None:
"""Failure-test role authorization."""
post_body = {
"database": "test_histograms",
"collection": "TEST",
"histogram": {
"Anything": True
},
}
with pytest.raises(requests.exceptions.HTTPError) as e:
db_rc.request_seq("POST", "/histogram", post_body)
assert e.response.status_code == 403 # Forbidden Error
db_rc.close()
async def main():
admin_token = Auth('secret').create_token('foo', payload={'role': 'admin'})
user_token = Auth('secret').create_token('foo', payload={'role': 'user'})
api = RestClient('http://localhost:8080/api', token=admin_token)
await api.request('POST', '/fruits', {'name': 'apple'})
await api.request('POST', '/fruits', {'name': 'banana'})
api = RestClient('http://localhost:8080/api', token=user_token)
ret = await api.request('GET', '/fruits')
if ret != {'apple': {'name': 'apple'}, 'banana': {'name': 'banana'}}:
print(ret)
print('FAIL')
else:
print('OK')
async def test_91_request_seq(requests_mock: Mock) -> None:
"""Test `request_seq()`."""
result = {"result": "the result"}
rpc = RestClient("http://test", "passkey", timeout=0.1)
def response(req: PreparedRequest, ctx: object) -> bytes: # pylint: disable=W0613
assert req.body is not None
_ = json_decode(req.body)
return json_encode(result).encode("utf-8")
requests_mock.post("/test", content=response)
ret = rpc.request_seq("POST", "test", {})
assert requests_mock.called
assert ret == result
def test_60_post_files_locations_1xN(self):
"""Test locations uniqueness under 1xN multiplicity."""
self.start_server()
token = self.get_token()
r = RestClient(self.address, token, timeout=1, retries=1)
# define the locations to be tested
loc1a = {'site': 'WIPAC', 'path': '/data/test/exp/IceCube/foo.dat'}
loc1b = {'site': 'DESY', 'path': '/data/test/exp/IceCube/foo.dat'}
loc1c = {'site': 'NERSC', 'path': '/data/test/exp/IceCube/foo.dat'}
loc1d = {'site': 'OSG', 'path': '/data/test/exp/IceCube/foo.dat'}
locs3a = [loc1a, loc1b, loc1c]
locs3b = [loc1b, loc1c, loc1d]
locs3c = [loc1a, loc1b, loc1d]
# define the files to be created
metadata = {
'logical_name': '/blah/data/exp/IceCube/blah.dat',
'checksum': {'sha512': hex('foo bar')},
'file_size': 1,
u'locations': [loc1a]
}
metadata2 = {
'logical_name': '/blah/data/exp/IceCube/blah2.dat',
'checksum': {'sha512': hex('foo bar2')},
'file_size': 2,
u'locations': locs3a
}
# create the first file; should be OK
data = r.request_seq('POST', '/api/files', metadata)
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('file', data)
url = data['file']
uid = url.split('/')[-1]
# check that the file was created properly
data = r.request_seq('GET', '/api/files')
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('files', data)
self.assertEqual(len(data['files']), 1)
self.assertTrue(any(uid == f['uuid'] for f in data['files']))
# create the second file; should NOT be OK
with self.assertRaises(Exception):
r.request_seq('POST', '/api/files', metadata2)
async def test_10_request(requests_mock: Mock) -> None:
"""Test `async request()`."""
result = {"result": "the result"}
rpc = RestClient("http://test", "passkey", timeout=0.1)
def response(req: PreparedRequest, ctx: object) -> bytes: # pylint: disable=W0613
assert req.body is not None
_ = json_decode(req.body)
return json_encode(result).encode("utf-8")
requests_mock.post("/test", content=response)
ret = await rpc.request("POST", "test", {})
assert requests_mock.called
auth_parts = requests_mock.last_request.headers['Authorization'].split(
' ', 1)
assert auth_parts[0].lower() == 'bearer'
assert auth_parts[1] == 'passkey'
assert ret == result
result2 = {"result2": "the result 2"}
def response2(req: PreparedRequest, ctx: object) -> bytes: # pylint: disable=W0613
assert req.body is not None
_ = json_decode(req.body)
return json_encode(result2).encode("utf-8")
requests_mock.post("/test2", content=response2)
ret = await rpc.request("POST", "/test2")
assert requests_mock.called
assert ret == result2
async def _do_work_claim(self) -> bool:
"""Claim a bundle and perform work on it."""
# 1. Ask the LTA DB for the next Bundle to be transferred
# configure a RestClient to talk to the LTA DB
lta_rc = RestClient(self.lta_rest_url,
token=self.lta_rest_token,
timeout=self.work_timeout_seconds,
retries=self.work_retries)
self.logger.info("Asking the LTA DB for a Bundle to transfer.")
pop_body = {"claimant": f"{self.name}-{self.instance_uuid}"}
response = await lta_rc.request(
'POST',
f'/Bundles/actions/pop?source={self.source_site}&dest={self.dest_site}&status={self.input_status}',
pop_body)
self.logger.info(f"LTA DB responded with: {response}")
bundle = response["bundle"]
if not bundle:
self.logger.info(
"LTA DB did not provide a Bundle to transfer. Going on vacation."
)
return False
# process the Bundle that we were given
try:
await self._replicate_bundle_to_destination_site(lta_rc, bundle)
except Exception as e:
await self._quarantine_bundle(lta_rc, bundle, f"{e}")
return False
# if we were successful at processing work, let the caller know
return True
async def client(username='******', groups=[], timeout=10):
await krs.users.create_user(username,
'first',
'last',
username + '@test',
rest_client=keycloak_bootstrap)
await krs.users.set_user_password(username,
'test',
rest_client=keycloak_bootstrap)
for group in groups:
await krs.groups.create_group(group,
rest_client=keycloak_bootstrap)
await krs.groups.add_user_group(group,
username,
rest_client=keycloak_bootstrap)
token = krs.apps.get_public_token(
username=username,
password='******',
scopes=['profile'],
client='user_mgmt',
openid_url=os.environ["AUTH_OPENID_URL"],
raw=True)
print(token)
return RestClient(f'http://localhost:{port}',
token=token,
timeout=timeout,
retries=0)
async def _do_work_claim(self) -> bool:
"""Claim a bundle and perform work on it."""
# 1. Ask the LTA DB for the next Bundle to be deleted
# configure a RestClient to talk to the LTA DB
lta_rc = RestClient(self.lta_rest_url,
token=self.lta_rest_token,
timeout=self.work_timeout_seconds,
retries=self.work_retries)
self.logger.info(
"Asking the LTA DB for a Bundle to check for TransferRequest being finished."
)
pop_body = {"claimant": f"{self.name}-{self.instance_uuid}"}
response = await lta_rc.request(
'POST',
f'/Bundles/actions/pop?source={self.source_site}&dest={self.dest_site}&status={self.input_status}',
pop_body)
self.logger.info(f"LTA DB responded with: {response}")
bundle = response["bundle"]
if not bundle:
self.logger.info(
"LTA DB did not provide a Bundle to check. Going on vacation.")
return False
# update the TransferRequest that spawned the Bundle, if necessary
await self._update_transfer_request(lta_rc, bundle)
# even if we processed a Bundle, take a break between Bundles
return False
async def patch_status_heartbeat(component: Component) -> bool:
"""PATCH /status/{component} to update LTA with a status heartbeat."""
component.logger.info("Sending status heartbeat")
# determine which resource to PATCH
status_route = f"/status/{component.type}"
status_url = urljoin(component.lta_rest_url, status_route)
# determine the body to PATCH with
status_body = {
component.name: {
"timestamp": datetime.utcnow().isoformat(),
"last_work_begin_timestamp": component.last_work_begin_timestamp,
"last_work_end_timestamp": component.last_work_end_timestamp,
}
}
# ask the base class to annotate the status body
status_update = component._do_status()
status_body[component.name].update(status_update)
# attempt to PATCH the status resource
component.logger.info(f"PATCH {status_url} - {status_body}")
try:
rc = RestClient(component.lta_rest_url,
token=component.lta_rest_token,
timeout=component.heartbeat_patch_timeout_seconds,
retries=component.heartbeat_patch_retries)
# Use the RestClient to PATCH our heartbeat to the LTA DB
await rc.request("PATCH", status_route, status_body)
except Exception as e:
# if there was a problem, yo I'll solve it
component.logger.error(f"Error trying to PATCH {status_route} with heartbeat")
component.logger.error(f"Error was: '{e}'", exc_info=True)
return False
# indicate to the caller that the heartbeat was successful
return True
async def _add_location_to_file_catalog(self, bundle_file: Dict[str, Any],
dest_path: str) -> bool:
"""Update File Catalog record with new Data Warehouse location."""
# configure a RestClient to talk to the File Catalog
fc_rc = RestClient(self.file_catalog_rest_url,
token=self.file_catalog_rest_token,
timeout=self.work_timeout_seconds,
retries=self.work_retries)
# extract the right variables from the metadata structure
fc_path = dest_path
fc_uuid = bundle_file["uuid"]
# add the new location to the File Catalog
new_location = {
"locations": [{
"site": "WIPAC",
"path": f"{fc_path}",
}]
}
self.logger.info(
f"POST /api/files/{fc_uuid}/locations - {new_location}")
# POST /api/files/{uuid}/locations will de-dupe locations for us
await fc_rc.request("POST", f"/api/files/{fc_uuid}/locations",
new_location)
# indicate that our file catalog updates were successful
return True
async def clear_catalog():
# configure a RestClient from the environment
config = from_environment(EXPECTED_CONFIG)
rc = RestClient(config["FILE_CATALOG_REST_URL"], token=config["FILE_CATALOG_REST_TOKEN"])
# while there are still files
clearing = True
while clearing:
try:
# get a list of up to 50 files
response = await rc.request("GET", "/api/files?start=0&limit=50")
files = response["files"]
# for each file that we found
for x in files:
# remove it from the file catalog
uuid = x["uuid"]
logical_name = x["logical_name"]
print(f"DELETE /api/files/{uuid} - {logical_name}")
response2 = await rc.request("DELETE", f"/api/files/{uuid}")
# if we didn't get any files back, we're done
if len(files) < 1:
clearing = False
except Exception as e:
# whoopsy daisy...
clearing = False
print(e)
async def clear_lta_transfer_requests():
# configure a RestClient from the environment
config = from_environment(EXPECTED_CONFIG)
rc = RestClient(config["LTA_REST_URL"], token=config["LTA_REST_TOKEN"])
# while there are still transfer requests
clearing = True
while clearing:
try:
# get a list of up to 50 transfer requests
# technically a lie; the LTA DB honors neither start nor limit
response = await rc.request("GET", "/TransferRequests?start=0&limit=50")
results = response["results"]
# for each file that we found
for x in results:
# remove it from the file catalog
uuid = x["uuid"]
print(f"DELETE /TransferRequests/{uuid}")
response2 = await rc.request("DELETE", f"/TransferRequests/{uuid}")
# if we didn't get any files back, we're done
if len(results) < 1:
clearing = False
except Exception as e:
# whoopsy daisy...
clearing = False
print(e)
async def test_22_request(requests_mock: Mock) -> None:
"""Test `async request()`."""
rpc = RestClient("http://test", "passkey", timeout=0.1)
requests_mock.get("/test", content=b'{"foo"}')
with pytest.raises(Exception):
_ = await rpc.request("GET", "test", {})
async def test_20_timeout(requests_mock: Mock) -> None:
"""Test timeout in `async request()`."""
rpc = RestClient("http://test", "passkey", timeout=0.1, backoff=False)
requests_mock.post("/test", exc=Timeout)
with pytest.raises(Timeout):
_ = await rpc.request("POST", "test", {})
async def test_21_ssl_error(requests_mock: Mock) -> None:
"""Test ssl error in `async request()`."""
rpc = RestClient("http://test", "passkey", timeout=0.1, backoff=False)
requests_mock.post("/test", exc=SSLError)
with pytest.raises(SSLError):
_ = await rpc.request("POST", "test", {})
async def _do_work_claim(self) -> bool:
"""Claim a transfer request and perform work on it."""
# 1. Ask the LTA DB for the next TransferRequest to be picked
# configure a RestClient to talk to the LTA DB
lta_rc = RestClient(self.lta_rest_url,
token=self.lta_rest_token,
timeout=self.work_timeout_seconds,
retries=self.work_retries)
self.logger.info("Asking the LTA DB for a TransferRequest to work on.")
pop_body = {"claimant": f"{self.name}-{self.instance_uuid}"}
response = await lta_rc.request(
'POST',
f'/TransferRequests/actions/pop?source={self.source_site}&dest={self.dest_site}',
pop_body)
self.logger.info(f"LTA DB responded with: {response}")
tr = response["transfer_request"]
if not tr:
self.logger.info(
"LTA DB did not provide a TransferRequest to work on. Going on vacation."
)
return False
# process the TransferRequest that we were given
try:
await self._do_work_transfer_request(lta_rc, tr)
except Exception as e:
await self._quarantine_transfer_request(lta_rc, tr, f"{e}")
raise e
# if we were successful at processing work, let the caller know
return True
async def test_11_request(requests_mock: Mock) -> None:
"""Test request in `async request()`."""
rpc = RestClient("http://test", "passkey", timeout=0.1)
requests_mock.get("/test", content=b"")
ret = await rpc.request("GET", "test", {})
assert requests_mock.called
assert ret is None
def __init__(self, lta_rest_url: str, lta_rest_token: str,
monitoring_interval: str = '60', logger: Any = None) -> None:
"""Initialize a Monitor object."""
self.logger = logger if logger else logging
self.interval = int(monitoring_interval)
self.rest = RestClient(lta_rest_url, lta_rest_token,
timeout=self.interval//10, retries=1)
self.running = False
def test_58_patch_files_uuid_unique_locations(self):
"""Test that locations is unique when updating a file."""
self.start_server()
token = self.get_token()
r = RestClient(self.address, token, timeout=1, retries=1)
# define the files to be created
metadata = {
'logical_name': '/blah/data/exp/IceCube/blah.dat',
'checksum': {'sha512': hex('foo bar')},
'file_size': 1,
u'locations': [{u'site': u'WIPAC', u'path': u'/blah/data/exp/IceCube/blah.dat'}]
}
metadata2 = {
'logical_name': '/blah/data/exp/IceCube/blah2.dat',
'checksum': {'sha512': hex('foo bar')},
'file_size': 1,
u'locations': [{u'site': u'WIPAC', u'path': u'/blah/data/exp/IceCube/blah2.dat'}]
}
# this is a PATCH to metadata; steps on metadata2's location
patch1 = {
'logical_name': '/blah/data/exp/IceCube/blah.dat',
'checksum': {'sha512': hex('foo bar2')},
'file_size': 2,
u'locations': [{u'site': u'WIPAC', u'path': u'/blah/data/exp/IceCube/blah2.dat'}]
}
# create the first file; should be OK
data = r.request_seq('POST', '/api/files', metadata)
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('file', data)
url = data['file']
uid = url.split('/')[-1]
# create the second file; should be OK
data = r.request_seq('POST', '/api/files', metadata2)
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('file', data)
# try to update the first file with a patch; should NOT be OK
with self.assertRaises(Exception):
r.request_seq('PATCH', '/api/files/' + uuid, patch1)
def test_get(db_rc: RestClient) -> None:
"""Run some test queries."""
databases = db_rc.request_seq("GET", "/databases/names")
print(databases)
for db in databases["databases"]:
db_request_body = {"database": db}
collections = db_rc.request_seq("GET", "/collections/names",
db_request_body)
print(collections)
for coll in collections["collections"]:
coll_request_body = {"database": db, "collection": coll}
histograms = db_rc.request_seq(
"GET", "/collections/histograms/names", coll_request_body)
print(histograms)
for histo_name in histograms["histograms"]:
histo_request_body = {
"database": db,
"collection": coll,
"name": histo_name,
}
histo = db_rc.request_seq("GET", "/histogram",
histo_request_body)
print(histo)
filelist = db_rc.request_seq("GET", "/files/names",
coll_request_body)
print(filelist)
db_rc.close()
def __init__(self, iceprodv1_pass: str, iceprodv2_token: str):
if not iceprodv1_pass:
raise RuntimeError("Missing IceProd v1 DB password")
elif not iceprodv2_token:
raise RuntimeError("Missing IceProd v2 REST token")
self._iceprodv1_pass = iceprodv1_pass
self._iceprodv2_rc = RestClient(
"https://iceprod2-api.icecube.wisc.edu", iceprodv2_token)
async def test_30_request(requests_mock: Mock) -> None:
"""Test `async request()` with headers."""
rpc = RestClient("http://test", "passkey", timeout=0.1)
requests_mock.get("/test", content=b"")
ret = await rpc.request("GET", "test", {}, {'foo': 'bar'})
assert requests_mock.called
assert requests_mock.last_request.headers['foo'] == 'bar'
assert ret is None
def db_rc() -> RestClient:
"""Get database REST client."""
token_json = requests.get(
"http://localhost:8888/token?scope=maddash:production"
).json()
rc = RestClient(
"http://localhost:8080", token=token_json["access"], timeout=5, retries=0
)
return rc
def test_10_collections(self):
self.start_server()
token = self.get_token()
r = RestClient(self.address, token, timeout=1, retries=1)
metadata = {
'collection_name': 'blah',
'owner': 'foo',
}
data = r.request_seq('POST', '/api/collections', metadata)
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('collection', data)
url = data['collection']
uid = url.split('/')[-1]
data = r.request_seq('GET', '/api/collections')
self.assertIn('collections', data)
self.assertIn(uid,{row['uuid'] for row in data['collections']})
def test_21_collection_by_name(self):
self.start_server()
token = self.get_token()
r = RestClient(self.address, token, timeout=1, retries=1)
metadata = {
'collection_name': 'blah',
'owner': 'foo',
}
data = r.request_seq('POST', '/api/collections', metadata)
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('collection', data)
url = data['collection']
uid = url.split('/')[-1]
data = r.request_seq('GET', '/api/collections/blah')
for k in metadata:
self.assertIn(k, data)
self.assertEqual(metadata[k], data[k])
def create_simprod_dbms_rest_connection() -> RestClient:
"""Return REST Client connection object."""
token_request_url = urljoin(token_server_url, "token?scope=maddash:web")
token_json = requests.get(token_request_url).json()
rc = RestClient(dbms_server_url,
token=token_json["access"],
timeout=5,
retries=0)
return rc
def index(
paths: List[str],
blacklist: List[str],
rest_client_args: RestClientArgs,
site: str,
indexer_flags: IndexerFlags,
) -> List[str]:
"""Index paths, excluding any matching the blacklist.
Return all child paths nested under any directories.
"""
if not isinstance(paths, list):
raise TypeError(f"`paths` object is not list {paths}")
if not paths:
return []
# Filter
paths = file_utils.sorted_unique_filepaths(list_of_filepaths=paths)
paths = [p for p in paths if not path_in_blacklist(p, blacklist)]
# Prep
fc_rc = RestClient(
rest_client_args["url"],
token=rest_client_args["token"],
timeout=rest_client_args["timeout"],
retries=rest_client_args["retries"],
)
manager = MetadataManager(
site,
basic_only=indexer_flags["basic_only"],
iceprodv2_rc_token=indexer_flags["iceprodv2_rc_token"],
iceprodv1_db_pass=indexer_flags["iceprodv1_db_pass"],
)
# Index
child_paths = asyncio.get_event_loop().run_until_complete(
index_paths(paths, manager, fc_rc, indexer_flags["patch"],
indexer_flags["dryrun"]))
fc_rc.close()
return child_paths
def test_50_post_files_unique_logical_name(self):
"""Test that logical_name is unique when creating a new file."""
self.start_server()
token = self.get_token()
r = RestClient(self.address, token, timeout=1, retries=1)
# define the file to be created
metadata = {
'logical_name': '/blah/data/exp/IceCube/blah.dat',
'checksum': {'sha512': hex('foo bar')},
'file_size': 1,
u'locations': [{u'site': u'WIPAC', u'path': u'/blah/data/exp/IceCube/blah.dat'}]
}
# create the file the first time; should be OK
data = r.request_seq('POST', '/api/files', metadata)
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('file', data)
url = data['file']
uid = url.split('/')[-1]
# check that the file was created properly
data = r.request_seq('GET', '/api/files')
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('files', data)
self.assertEqual(len(data['files']), 1)
self.assertTrue(any(uid == f['uuid'] for f in data['files']))
# create the file the second time; should NOT be OK
with self.assertRaises(Exception):
data = r.request_seq('POST', '/api/files', metadata)
# check that the second file was not created
data = r.request_seq('GET', '/api/files')
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('files', data)
self.assertEqual(len(data['files']), 1)
self.assertTrue(any(uid == f['uuid'] for f in data['files']))
def test_102_request_stream() -> None:
"""Test `request_stream()` where there's only one response."""
mock_url = "http://test"
rpc = RestClient(mock_url, "passkey", timeout=1)
one_liners = [
[b'"with-a-newline"\n'],
[b'"w/o-a-newline"'],
]
for expected_stream in one_liners:
json_stream = [j for t in expected_stream
if (j := _jsonify(t))] # no blanks
async def _add_bundle_to_file_catalog(self, lta_rc: RestClient,
bundle: BundleType) -> bool:
"""Add a FileCatalog entry for the bundle, then update existing records."""
# configure a RestClient to talk to the File Catalog
fc_rc = RestClient(self.file_catalog_rest_url,
token=self.file_catalog_rest_token,
timeout=self.work_timeout_seconds,
retries=self.work_retries)
# determine the path where the bundle is stored on hpss
data_warehouse_path = bundle["path"]
basename = os.path.basename(bundle["bundle_path"])
stupid_python_path = os.path.sep.join(
[self.tape_base_path, data_warehouse_path, basename])
hpss_path = os.path.normpath(stupid_python_path)
# create a File Catalog entry for the bundle itself
bundle_uuid = bundle["uuid"]
right_now = now()
file_record = {
"uuid":
bundle_uuid,
"logical_name":
hpss_path,
"checksum":
bundle["checksum"],
"locations": [{
"site": "NERSC",
"path": hpss_path,
"hpss": True,
"online": False,
}],
"file_size":
bundle["size"],
"lta": {
"date_archived": right_now,
},
}
# add the bundle file to the File Catalog
try:
self.logger.info(f"POST /api/files - {hpss_path}")
await fc_rc.request("POST", "/api/files", file_record)
except Exception as e:
self.logger.error(f"Error: POST /api/files - {hpss_path}")
self.logger.error(f"Message: {e}")
bundle_uuid = bundle["uuid"]
self.logger.info(f"PATCH /api/files/{bundle_uuid}")
await fc_rc.request("PATCH", f"/api/files/{bundle_uuid}",
file_record)
# update the File Catalog for each file contained in the bundle
await self._update_files_in_file_catalog(fc_rc, lta_rc, bundle,
hpss_path)
# indicate that our file catalog updates were successful
return True
def client(role='read', timeout=0.1):
if CONFIG['AUTH_ISSUER']:
r = requests.get(CONFIG['AUTH_ISSUER'] + '/token',
params={'scope': f'prometheus-reconfig:{role}'})
r.raise_for_status()
t = r.json()['access']
else:
raise Exception('testing token service not defined')
print(t)
return RestClient(f'http://localhost:{http_server_port}',
token=t,
timeout=timeout,
retries=0)
def test_30_archive(self):
self.start_server()
token = self.get_token()
r = RestClient(self.address, token, timeout=1, retries=1)
metadata = {
u'logical_name': u'blah',
u'checksum': {u'sha512':hex('foo bar')},
u'file_size': 1,
u'locations': [{u'site':u'test',u'path':u'blah.dat'}]
}
metadata2 = {
u'logical_name': u'blah2',
u'checksum': {u'sha512':hex('foo bar baz')},
u'file_size': 2,
u'locations': [{u'site':u'test',u'path':u'blah.dat',u'archive':True}]
}
data = r.request_seq('POST', '/api/files', metadata)
url = data['file']
uid = url.split('/')[-1]
data = r.request_seq('POST', '/api/files', metadata2)
url2 = data['file']
uid2 = url2.split('/')[-1]
data = r.request_seq('GET', '/api/files')
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('files', data)
self.assertEqual(len(data['files']), 1)
self.assertTrue(any(uid == f['uuid'] for f in data['files']))
self.assertFalse(any(uid2 == f['uuid'] for f in data['files']))
data = r.request_seq('GET', '/api/files', {'query':json_encode({'locations.archive':True})})
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('files', data)
self.assertEqual(len(data['files']), 1)
self.assertFalse(any(uid == f['uuid'] for f in data['files']))
self.assertTrue(any(uid2 == f['uuid'] for f in data['files']))
def test_30_collection_files(self):
self.start_server()
token = self.get_token()
r = RestClient(self.address, token, timeout=1, retries=1)
metadata = {
'collection_name': 'blah',
'owner': 'foo',
}
data = r.request_seq('POST', '/api/collections', metadata)
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('collection', data)
url = data['collection']
uid = url.split('/')[-1]
data = r.request_seq('GET', '/api/collections/blah/files')
self.assertEqual(data['files'], [])
# add a file
metadata = {
'logical_name': 'blah',
'checksum': {'sha512':hex('foo bar')},
'file_size': 1,
u'locations': [{u'site':u'test',u'path':u'blah.dat'}]
}
data = r.request_seq('POST', '/api/files', metadata)
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('file', data)
url = data['file']
uid = url.split('/')[-1]
data = r.request_seq('GET', '/api/collections/blah/files',
{'keys':'uuid|logical_name|checksum|locations'})
self.assertEqual(len(data['files']), 1)
self.assertEqual(data['files'][0]['uuid'], uid)
self.assertEqual(data['files'][0]['checksum'], metadata['checksum'])
def main():
parser = argparse.ArgumentParser(description='manually run IceProd i3exec')
parser.add_argument('-t', '--token',help='auth token')
parser.add_argument('-d','--dataset',type=int,help='dataset number')
parser.add_argument('-j','--job',type=int,help='job number (optional)')
args = parser.parse_args()
args = vars(args)
logging.basicConfig(level=logging.DEBUG)
rpc = RestClient('https://iceprod2-api.icecube.wisc.edu', args['token'])
datasets = rpc.request_seq('GET', '/datasets', {'keys': 'dataset_id|dataset'})
dataset_id = None
for d in datasets:
if datasets[d]['dataset'] == args['dataset']:
dataset_id = d
break
else:
raise Exception('bad dataset num')
dataset = rpc.request_seq('GET', f'/datasets/{dataset_id}')
config = rpc.request_seq('GET', f'/config/{dataset_id}')
jobs = rpc.request_seq('GET', f'/datasets/{dataset_id}/jobs', {'status': 'processing|errors'})
if args['job']:
jobs = {j:jobs[j] for j in jobs if jobs[j]['job_index'] == args['job']}
if not jobs:
raise Exception('no jobs found')
for job_id in jobs:
tasks = rpc.request_seq('GET', f'/datasets/{dataset_id}/tasks',
{'job_id': job_id, 'keys': 'task_id|task_index|name|depends',
'status': 'waiting|queued|reset|failed'})
for task_id in sorted(tasks, key=lambda t:tasks[t]['task_index']):
print(f'processing {dataset["dataset"]} {jobs[job_id]["job_index"]} {tasks[task_id]["name"]}')
write_config(config, 'config.json', dataset_id, args['dataset'], task_id)
run(token=args['token'], config='config.json',
jobs_submitted=dataset['jobs_submitted'],
job=jobs[job_id]['job_index'],
task=tasks[task_id]['name'])
def test_20_file(self):
self.start_server()
token = self.get_token()
r = RestClient(self.address, token, timeout=1, retries=1)
metadata = {
u'logical_name': u'blah',
u'checksum': {u'sha512':hex('foo bar')},
u'file_size': 1,
u'locations': [{u'site':u'test',u'path':u'blah.dat'}]
}
data = r.request_seq('POST', '/api/files', metadata)
url = data['file']
data = r.request_seq('GET', url)
data.pop('_links')
data.pop('meta_modify_date')
data.pop('uuid')
self.assertDictEqual(metadata, data)
metadata['test'] = 100
metadata_cpy = metadata.copy()
metadata_cpy['uuid'] = 'something else'
with self.assertRaises(Exception):
data = r.request_seq('PUT', url, metadata_cpy)
data = r.request_seq('PUT', url, metadata)
data.pop('_links')
data.pop('meta_modify_date')
data.pop('uuid')
self.assertDictEqual(metadata, data)
data = r.request_seq('GET', url)
data.pop('_links')
data.pop('meta_modify_date')
data.pop('uuid')
self.assertDictEqual(metadata, data)
metadata['test2'] = 200
data = r.request_seq('PATCH', url, {'test2':200})
data.pop('_links')
data.pop('meta_modify_date')
data.pop('uuid')
self.assertDictEqual(metadata, data)
data = r.request_seq('GET', url)
data.pop('_links')
data.pop('meta_modify_date')
data.pop('uuid')
self.assertDictEqual(metadata, data)
data = r.request_seq('DELETE', url)
# second delete should raise error
with self.assertRaises(Exception):
data = r.request_seq('DELETE', url)
with self.assertRaises(Exception):
data = r.request_seq('POST', url)
class ServerComms:
"""
Setup JSONRPC communications with the IceProd server.
Args:
url (str): address to connect to
passkey (str): passkey for authorization/authentication
config (:py:class:`iceprod.server.exe.Config`): Config object
**kwargs: passed to JSONRPC
"""
def __init__(self, url, passkey, config, **kwargs):
self.rest = RestClient(address=url,token=passkey,**kwargs)
async def download_task(self, gridspec, resources={}):
"""
Download new task(s) from the server.
Args:
gridspec (str): gridspec the pilot was submitted from
resources (dict): resources available in the pilot
Returns:
list: list of task configs
"""
hostname = functions.gethostname()
domain = '.'.join(hostname.split('.')[-2:])
try:
ifaces = functions.getInterfaces()
except Exception:
ifaces = None
resources = deepcopy(resources)
if 'gpu' in resources and isinstance(resources['gpu'],list):
resources['gpu'] = len(resources['gpu'])
os_type = os.environ['OS_ARCH'] if 'OS_ARCH' in os.environ else None
if os_type:
resources['os'] = os_type
task = await self.rest.request('POST', '/task_actions/process',
{'gridspec': gridspec,
'hostname': hostname,
'domain': domain,
'ifaces': ifaces,
'requirements': resources,
})
if not task:
return None
# get config
try:
config = await self.rest.request('GET', '/config/{}'.format(task['dataset_id']))
if not isinstance(config, dataclasses.Job):
config = dict_to_dataclasses(config)
except Exception:
logging.warning('failed to get dataset config for dataset %s', task['dataset_id'])
await self.task_kill(task['task_id'], dataset_id=task['dataset_id'],
reason='failed to download dataset config')
raise
# fill in options
if 'options' not in config:
config['options'] = {}
config['options']['task_id'] = task['task_id']
config['options']['job_id'] = task['job_id']
config['options']['dataset_id'] = task['dataset_id']
config['options']['task'] = task['task_index']
if 'requirements' in task:
config['options']['resources'] = {k:task['requirements'][k] for k in Resources.defaults}
try:
job = await self.rest.request('GET', '/jobs/{}'.format(task['job_id']))
config['options']['job'] = job['job_index']
except Exception:
logging.warning('failed to get job %s', task['job_id'])
await self.task_kill(task['task_id'], dataset_id=task['dataset_id'],
reason='failed to download job')
raise
try:
dataset = await self.rest.request('GET', '/datasets/{}'.format(task['dataset_id']))
config['options']['dataset'] = dataset['dataset']
config['options']['jobs_submitted'] = dataset['jobs_submitted']
config['options']['tasks_submitted'] = dataset['tasks_submitted']
config['options']['debug'] = dataset['debug']
except Exception:
logging.warning('failed to get dataset %s', task['dataset_id'])
await self.task_kill(task['task_id'], dataset_id=task['dataset_id'],
reason='failed to download dataset')
raise
return [config]
async def task_files(self, dataset_id, task_id):
"""
Get the task files for a dataset and task.
Args:
dataset_id (str): dataset_id
task_id (str): task_id
Returns:
list: list of :py:class:`iceprod.core.dataclasses.Data` objects
"""
ret = await self.rest.request('GET', '/datasets/{}/task_files/{}'.format(dataset_id, task_id))
data = []
for r in ret['files']:
d = dataclasses.Data(r)
if not d.valid():
raise Exception('returned Data not valid')
data.append(d)
return data
async def processing(self, task_id):
"""
Tell the server that we are processing this task.
Only used for single task config, not for pilots.
Args:
task_id (str): task_id to mark as processing
"""
await self.rest.request('PUT', '/tasks/{}/status'.format(task_id),
{'status': 'processing'})
async def finish_task(self, task_id, dataset_id=None, stats={},
stat_filter=None, start_time=None, resources=None):
"""
Finish a task.
Args:
task_id (str): task_id of task
dataset_id (str): (optional) dataset_id of task
stats (dict): (optional) task statistics
stat_filter (iterable): (optional) stat filter by keywords
start_time (float): (optional) task start time in unix seconds
resources (dict): (optional) task resource usage
"""
if stat_filter:
# filter task stats
stats = {k:stats[k] for k in stats if k in stat_filter}
hostname = functions.gethostname()
domain = '.'.join(hostname.split('.')[-2:])
if start_time:
t = time.time() - start_time
elif resources and 'time' in resources and resources['time']:
t = int(resources['time']*3600)
else:
t = None
iceprod_stats = {
'hostname': hostname,
'domain': domain,
'time_used': t,
'task_stats': stats,
'time': datetime.utcnow().isoformat(),
}
if resources:
iceprod_stats['resources'] = resources
if dataset_id:
iceprod_stats['dataset_id'] = dataset_id
await self.rest.request('POST', '/tasks/{}/task_stats'.format(task_id),
iceprod_stats)
data = {}
if t:
data['time_used'] = t
await self.rest.request('POST', '/tasks/{}/task_actions/complete'.format(task_id), data)
async def still_running(self, task_id):
"""
Check if the task should still be running according to the DB.
Args:
task_id (str): task_id of task
"""
ret = await self.rest.request('GET', '/tasks/{}'.format(task_id))
if (not ret) or 'status' not in ret or ret['status'] != 'processing':
raise Exception('task should be stopped')
async def task_error(self, task_id, dataset_id=None, stats={}, start_time=None, reason=None, resources=None):
"""
Tell the server about the error experienced
Args:
task_id (str): task_id of task
dataset_id (str): (optional) dataset_id of task
stats (dict): (optional) task statistics
start_time (float): (optional) task start time in unix seconds
reason (str): (optional) one-line summary of error
resources (dict): (optional) task resource usage
"""
iceprod_stats = {}
try:
hostname = functions.gethostname()
domain = '.'.join(hostname.split('.')[-2:])
if start_time:
t = time.time() - start_time
elif resources and 'time' in resources and resources['time']:
t = int(resources['time']*3600)
else:
t = None
iceprod_stats = {
'task_id': task_id,
'hostname': hostname,
'domain': domain,
'time_used': t,
'task_stats': json.dumps(stats),
'time': datetime.utcnow().isoformat(),
'error_summary': reason if reason else '',
}
if dataset_id:
iceprod_stats['dataset_id'] = dataset_id
if resources:
iceprod_stats['resources'] = resources
except Exception:
logging.warning('failed to collect error info', exc_info=True)
try:
await self.rest.request('POST', '/tasks/{}/task_stats'.format(task_id),
iceprod_stats)
except Exception:
logging.warning('failed to post task_stats for %r', task_id, exc_info=True)
data = {}
if t:
data['time_used'] = t
if resources:
data['resources'] = resources
if reason:
data['reason'] = reason
await self.rest.request('POST', '/tasks/{}/task_actions/reset'.format(task_id), data)
async def task_kill(self, task_id, dataset_id=None, resources=None, reason=None, message=None):
"""
Tell the server that we killed a task.
Args:
task_id (str): task_id of task
dataset_id (str): (optional) dataset_id of task
resources (dict): (optional) used resources
reason (str): (optional) short summary for kill
message (str): (optional) long message to replace log upload
"""
if not reason:
reason = 'killed'
if not message:
message = reason
try:
hostname = functions.gethostname()
domain = '.'.join(hostname.split('.')[-2:])
iceprod_stats = {
'task_id': task_id,
'hostname': hostname,
'domain': domain,
'time': datetime.utcnow().isoformat(),
'error_summary': reason if reason else '',
}
if dataset_id:
iceprod_stats['dataset_id'] = dataset_id
if resources:
iceprod_stats['resources'] = resources
except Exception:
logging.warning('failed to collect error info', exc_info=True)
iceprod_stats = {}
try:
await self.rest.request('POST', '/tasks/{}/task_stats'.format(task_id),
iceprod_stats)
except Exception:
logging.warning('failed to post task_stats for %r', task_id, exc_info=True)
data = {}
if resources and 'time' in resources and resources['time']:
data['time_used'] = resources['time']*3600.
if resources:
data['resources'] = resources
if reason:
data['reason'] = reason
else:
data['data'] = 'task killed'
await self.rest.request('POST', '/tasks/{}/task_actions/reset'.format(task_id), data)
data = {'name': 'stdlog', 'task_id': task_id}
if dataset_id:
data['dataset_id'] = dataset_id
if message:
data['data'] = message
elif reason:
data['data'] = reason
else:
data['data'] = 'task killed'
await self.rest.request('POST', '/logs', data)
data.update({'name':'stdout', 'data': ''})
await self.rest.request('POST', '/logs', data)
data.update({'name':'stderr', 'data': ''})
await self.rest.request('POST', '/logs', data)
async def _upload_logfile(self, name, filename, task_id=None, dataset_id=None):
"""Upload a log file"""
data = {'name': name}
if task_id:
data['task_id'] = task_id
if dataset_id:
data['dataset_id'] = dataset_id
try:
with open(filename) as f:
data['data'] = f.read()
except Exception as e:
data['data'] = str(e)
await self.rest.request('POST', '/logs', data)
async def uploadLog(self, **kwargs):
"""Upload log file"""
logging.getLogger().handlers[0].flush()
await self._upload_logfile('stdlog', os.path.abspath(constants['stdlog']), **kwargs)
async def uploadErr(self, filename=None, **kwargs):
"""Upload stderr file"""
if not filename:
sys.stderr.flush()
filename = os.path.abspath(constants['stderr'])
await self._upload_logfile('stderr', filename, **kwargs)
async def uploadOut(self, filename=None, **kwargs):
"""Upload stdout file"""
if not filename:
sys.stdout.flush()
filename = os.path.abspath(constants['stdout'])
await self._upload_logfile('stdout', filename, **kwargs)
async def create_pilot(self, **kwargs):
"""
Create an entry in the pilot table.
Args:
**kwargs: passed through to rest function
Returns:
str: pilot id
"""
ret = await self.rest.request('POST', '/pilots', kwargs)
return ret['result']
async def update_pilot(self, pilot_id, **kwargs):
"""
Update the pilot table.
Args:
pilot_id (str): pilot id
**kwargs: passed through to rest function
"""
await self.rest.request('PATCH', '/pilots/{}'.format(pilot_id), kwargs)
async def delete_pilot(self, pilot_id, **kwargs):
"""
Delete the pilot.
Args:
pilot_id (str): pilot id
"""
await self.rest.request('DELETE', '/pilots/{}'.format(pilot_id))
# --- synchronous versions to be used from a signal handler
# --- or other non-async code
def task_kill_sync(self, task_id, dataset_id=None, resources=None, reason=None, message=None):
"""
Tell the server that we killed a task (synchronous version).
Args:
task_id (str): task_id of task
dataset_id (str): (optional) dataset_id of task
resources (dict): (optional) used resources
reason (str): (optional) short summary for kill
message (str): (optional) long message to replace log upload
"""
if not reason:
reason = 'killed'
if not message:
message = reason
try:
hostname = functions.gethostname()
domain = '.'.join(hostname.split('.')[-2:])
iceprod_stats = {
'task_id': task_id,
'hostname': hostname,
'domain': domain,
'time': datetime.utcnow().isoformat(),
'error_summary': reason if reason else '',
}
if dataset_id:
iceprod_stats['dataset_id'] = dataset_id
if resources:
iceprod_stats['resources'] = resources
except Exception:
logging.warning('failed to collect error info', exc_info=True)
iceprod_stats = {}
try:
self.rest.request_seq('POST', '/tasks/{}/task_stats'.format(task_id),
iceprod_stats)
except Exception:
logging.warning('failed to post task_stats for %r', task_id, exc_info=True)
data = {}
if resources and 'time' in resources and resources['time']:
data['time_used'] = resources['time']*3600.
if resources:
data['resources'] = resources
if reason:
data['reason'] = reason
else:
data['data'] = 'task killed'
self.rest.request_seq('POST', '/tasks/{}/task_actions/reset'.format(task_id), data)
data = {'name': 'stdlog', 'task_id': task_id}
if dataset_id:
data['dataset_id'] = dataset_id
if message:
data['data'] = message
elif reason:
data['data'] = reason
else:
data['data'] = 'task killed'
self.rest.request_seq('POST', '/logs', data)
data.update({'name':'stdout', 'data': ''})
self.rest.request_seq('POST', '/logs', data)
data.update({'name':'stderr', 'data': ''})
self.rest.request_seq('POST', '/logs', data)
def update_pilot_sync(self, pilot_id, **kwargs):
"""
Update the pilot table (synchronous version).
Args:
pilot_id (str): pilot id
**kwargs: passed through to rpc function
"""
self.rest.request_seq('PATCH', '/pilots/{}'.format(pilot_id), kwargs)
def delete_pilot_sync(self, pilot_id, **kwargs):
"""
Delete the pilot (synchronous version).
Args:
pilot_id (str): pilot id
"""
self.rest.request_seq('DELETE', '/pilots/{}'.format(pilot_id))
def test_40_simple_query(self):
self.start_server()
token = self.get_token()
r = RestClient(self.address, token, timeout=1, retries=1)
metadata = {
u'logical_name': u'blah',
u'checksum': {u'sha512':hex('foo bar')},
u'file_size': 1,
u'locations': [{u'site':u'test',u'path':u'blah.dat'}],
u'processing_level':u'level2',
u'run_number':12345,
u'first_event':345,
u'last_event':456,
u'iceprod':{
u'dataset':23453,
},
u'offline':{
u'season':2017,
},
}
metadata2 = {
u'logical_name': u'blah2',
u'checksum': {u'sha512':hex('foo bar baz')},
u'file_size': 2,
u'locations': [{u'site':u'test',u'path':u'blah2.dat'}],
u'processing_level':u'level2',
r'run_number':12356,
u'first_event':578,
u'last_event':698,
u'iceprod':{
u'dataset':23454,
},
u'offline':{
u'season':2017,
},
}
data = r.request_seq('POST', '/api/files', metadata)
url = data['file']
uid = url.split('/')[-1]
data = r.request_seq('POST', '/api/files', metadata2)
url2 = data['file']
uid2 = url2.split('/')[-1]
data = r.request_seq('GET', '/api/files')
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('files', data)
self.assertEqual(len(data['files']), 2)
self.assertTrue(any(uid == f['uuid'] for f in data['files']))
self.assertTrue(any(uid2 == f['uuid'] for f in data['files']))
data = r.request_seq('GET', '/api/files', {'processing_level':'level2'})
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('files', data)
self.assertEqual(len(data['files']), 2)
self.assertTrue(any(uid == f['uuid'] for f in data['files']))
self.assertTrue(any(uid2 == f['uuid'] for f in data['files']))
data = r.request_seq('GET', '/api/files', {'run_number':12345})
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('files', data)
self.assertEqual(len(data['files']), 1)
self.assertTrue(any(uid == f['uuid'] for f in data['files']))
self.assertFalse(any(uid2 == f['uuid'] for f in data['files']))
data = r.request_seq('GET', '/api/files', {'dataset':23454})
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('files', data)
self.assertEqual(len(data['files']), 1)
self.assertFalse(any(uid == f['uuid'] for f in data['files']))
self.assertTrue(any(uid2 == f['uuid'] for f in data['files']))
data = r.request_seq('GET', '/api/files', {'event_id':400})
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('files', data)
self.assertEqual(len(data['files']), 1)
self.assertTrue(any(uid == f['uuid'] for f in data['files']))
self.assertFalse(any(uid2 == f['uuid'] for f in data['files']))
data = r.request_seq('GET', '/api/files', {'season':2017})
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('files', data)
self.assertEqual(len(data['files']), 2)
self.assertTrue(any(uid == f['uuid'] for f in data['files']))
self.assertTrue(any(uid2 == f['uuid'] for f in data['files']))
data = r.request_seq('GET', '/api/files', {'event_id':400, 'keys':'|'.join(['checksum','file_size','uuid'])})
self.assertIn('_links', data)
self.assertIn('self', data['_links'])
self.assertIn('files', data)
self.assertEqual(len(data['files']), 1)
self.assertTrue(any(uid == f['uuid'] for f in data['files']))
self.assertFalse(any(uid2 == f['uuid'] for f in data['files']))
self.assertIn('checksum', data['files'][0])
self.assertIn('file_size', data['files'][0])
def __init__(self, url, passkey, config, **kwargs):
self.rest = RestClient(address=url,token=passkey,**kwargs)
