FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

Merge pull request #697 from Kaggle/add-dataset-gcs · AI-For-Rural/docker-python@32a7b7c · GitHub

  1. AI-For-Rural
  2. docker-python

Repository navigation

Commit 32a7b7c

Browse files
mcollins42
authored
Merge pull request Kaggle#697 from Kaggle/add-dataset-gcs
Add a Kaggle dataset client to support getting GCS paths
2 parents 0f199da + 40fd3eb commit 32a7b7c

5 files changed

Lines changed: 214 additions & 0 deletions

File tree

‎Dockerfile‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -533,6 +533,8 @@ RUN echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] http://packages.c
533533
ENV PYTHONUSERBASE "/root/.local"
534534
ADD patches/kaggle_gcp.py /root/.local/lib/python3.6/site-packages/kaggle_gcp.py
535535
ADD patches/kaggle_secrets.py /root/.local/lib/python3.6/site-packages/kaggle_secrets.py
536+
ADD patches/kaggle_web_client.py /root/.local/lib/python3.6/site-packages/kaggle_web_client.py
537+
ADD patches/kaggle_datasets.py /root/.local/lib/python3.6/site-packages/kaggle_datasets.py
536538
ADD patches/log.py /root/.local/lib/python3.6/site-packages/log.py
537539
ADD patches/sitecustomize.py /root/.local/lib/python3.6/site-packages/sitecustomize.py
538540

‎patches/kaggle_datasets.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
import os
2+
from kaggle_web_client import KaggleWebClient
3+
4+
_KAGGLE_TPU_NAME_ENV_VAR_NAME = 'TPU_NAME'
5+
6+
class KaggleDatasets:
7+
GET_GCS_PATH_ENDPOINT = '/requests/CopyDatasetVersionToKnownGcsBucketRequest'
8+
9+
# Integration types for GCS
10+
AUTO_ML = 1
11+
TPU = 2
12+
13+
def __init__(self):
14+
self.web_client = KaggleWebClient()
15+
self.has_tpu = os.getenv(_KAGGLE_TPU_NAME_ENV_VAR_NAME) is not None
16+
17+
def get_gcs_path(self, dataset_dir: str = None) -> str:
18+
integration_type = self.TPU if self.has_tpu else self.AUTO_ML
19+
data = {
20+
'MountSlug': dataset_dir,
21+
'IntegrationType': integration_type,
22+
}
23+
result = self.web_client.make_post_request(data, self.GET_GCS_PATH_ENDPOINT)
24+
return result['destinationBucket']

‎patches/kaggle_secrets.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,8 @@ def __init__(self):
6767
self.headers = {'Content-type': 'application/json'}
6868

6969
def _make_post_request(self, data: dict, endpoint: str = GET_USER_SECRET_ENDPOINT) -> dict:
70+
# TODO(b/148309982) This code and the code in the constructor should be
71+
# removed and this class should use the new KaggleWebClient class instead.
7072
url = f'{self.url_base}{endpoint}'
7173
request_body = dict(data)
7274
request_body['JWE'] = self.jwt_token

‎patches/kaggle_web_client.py‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
import json
2+
import os
3+
import socket
4+
import urllib.request
5+
from datetime import datetime, timedelta
6+
from enum import Enum, unique
7+
from typing import Optional, Tuple
8+
from urllib.error import HTTPError, URLError
9+
from kaggle_secrets import (_KAGGLE_DEFAULT_URL_BASE,
10+
_KAGGLE_URL_BASE_ENV_VAR_NAME,
11+
_KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME,
12+
CredentialError, BackendError, ValidationError)
13+
14+
class KaggleWebClient:
15+
TIMEOUT_SECS = 90
16+
17+
def __init__(self):
18+
url_base_override = os.getenv(_KAGGLE_URL_BASE_ENV_VAR_NAME)
19+
self.url_base = url_base_override or _KAGGLE_DEFAULT_URL_BASE
20+
# Follow the OAuth 2.0 Authorization standard (https://tools.ietf.org/html/rfc6750)
21+
self.jwt_token = os.getenv(_KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME)
22+
if self.jwt_token is None:
23+
raise CredentialError(
24+
'A JWT Token is required to call Kaggle, '
25+
f'but none found in environment variable {_KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME}')
26+
self.headers = {
27+
'Content-type': 'application/json',
28+
'Authorization': f'Bearer {self.jwt_token}',
29+
}
30+
31+
def make_post_request(self, data: dict, endpoint: str) -> dict:
32+
url = f'{self.url_base}{endpoint}'
33+
request_body = dict(data)
34+
request_body['JWE'] = self.jwt_token
35+
req = urllib.request.Request(url, headers=self.headers, data=bytes(
36+
json.dumps(request_body), encoding="utf-8"))
37+
try:
38+
with urllib.request.urlopen(req, timeout=self.TIMEOUT_SECS) as response:
39+
response_json = json.loads(response.read())
40+
if not response_json.get('wasSuccessful') or 'result' not in response_json:
41+
raise BackendError(
42+
f'Unexpected response from the service. Response: {response_json}.')
43+
return response_json['result']
44+
except (URLError, socket.timeout) as e:
45+
if isinstance(
46+
e, socket.timeout) or isinstance(
47+
e.reason, socket.timeout):
48+
raise ConnectionError(
49+
'Timeout error trying to communicate with service. Please ensure internet is on.') from e
50+
raise ConnectionError(
51+
'Connection error trying to communicate with service.') from e
52+
except HTTPError as e:
53+
if e.code == 401 or e.code == 403:
54+
raise CredentialError(
55+
f'Service responded with error code {e.code}.'
56+
' Please ensure you have access to the resource.') from e
57+
raise BackendError('Unexpected response from the service.') from e

‎tests/test_datasets.py‎

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
import json
2+
import os
3+
import threading
4+
import unittest
5+
from http.server import BaseHTTPRequestHandler, HTTPServer
6+
from test.support import EnvironmentVarGuard
7+
from urllib.parse import urlparse
8+
9+
from kaggle_secrets import (_KAGGLE_URL_BASE_ENV_VAR_NAME,
10+
_KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME,
11+
CredentialError, BackendError, ValidationError)
12+
from kaggle_web_client import KaggleWebClient
13+
from kaggle_datasets import KaggleDatasets, _KAGGLE_TPU_NAME_ENV_VAR_NAME
14+
15+
_TEST_JWT = 'test-secrets-key'
16+
17+
_TPU_GCS_BUCKET = 'gs://kds-tpu-ea1971a458ffd4cd51389e7574c022ecc0a82bb1b52ccef08c8a'
18+
_AUTOML_GCS_BUCKET = 'gs://kds-automl-ea1971a458ffd4cd51389e7574c022ecc0a82bb1b52ccef08c8a'
19+
20+
class GcsDatasetsHTTPHandler(BaseHTTPRequestHandler):
21+
22+
def set_request(self):
23+
raise NotImplementedError()
24+
25+
def get_response(self):
26+
raise NotImplementedError()
27+
28+
def do_HEAD(s):
29+
s.send_response(200)
30+
31+
def do_POST(s):
32+
s.set_request()
33+
s.send_response(200)
34+
s.send_header("Content-type", "application/json")
35+
s.end_headers()
36+
s.wfile.write(json.dumps(s.get_response()).encode("utf-8"))
37+
38+
39+
class TestDatasets(unittest.TestCase):
40+
SERVER_ADDRESS = urlparse(os.getenv(_KAGGLE_URL_BASE_ENV_VAR_NAME, default="http://127.0.0.1:8001"))
41+
42+
def _test_client(self, client_func, expected_path, expected_body, is_tpu=True, success=True):
43+
_request = {}
44+
45+
class GetGcsPathHandler(GcsDatasetsHTTPHandler):
46+
47+
def set_request(self):
48+
_request['path'] = self.path
49+
content_len = int(self.headers.get('Content-Length'))
50+
_request['body'] = json.loads(self.rfile.read(content_len))
51+
_request['headers'] = self.headers
52+
53+
def get_response(self):
54+
if success:
55+
gcs_path = _TPU_GCS_BUCKET if is_tpu else _AUTOML_GCS_BUCKET
56+
return {'result': {
57+
'destinationBucket': gcs_path,
58+
'destinationPath': None}, 'wasSuccessful': "true"}
59+
else:
60+
return {'wasSuccessful': "false"}
61+
62+
env = EnvironmentVarGuard()
63+
env.set(_KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME, _TEST_JWT)
64+
if is_tpu:
65+
env.set(_KAGGLE_TPU_NAME_ENV_VAR_NAME, 'FAKE_TPU')
66+
with env:
67+
with HTTPServer((self.SERVER_ADDRESS.hostname, self.SERVER_ADDRESS.port), GetGcsPathHandler) as httpd:
68+
threading.Thread(target=httpd.serve_forever).start()
69+
70+
try:
71+
client_func()
72+
finally:
73+
httpd.shutdown()
74+
75+
path, headers, body = _request['path'], _request['headers'], _request['body']
76+
self.assertEqual(
77+
path,
78+
expected_path,
79+
msg="Fake server did not receive the right request from the KaggleDatasets client.")
80+
self.assertEqual(
81+
body,
82+
expected_body,
83+
msg="Fake server did not receive the right body from the KaggleDatasets client.")
84+
self.assertIn('Content-Type', headers.keys(),
85+
msg="Fake server did not receive a Content-Type header from the KaggleDatasets client.")
86+
self.assertEqual('application/json', headers.get('Content-Type'),
87+
msg="Fake server did not receive an application/json content type header from the KaggleDatasets client.")
88+
self.assertIn('Authorization', headers.keys(),
89+
msg="Fake server did not receive an Authorization header from the KaggleDatasets client.")
90+
self.assertEqual(f'Bearer {_TEST_JWT}', headers.get('Authorization'),
91+
msg="Fake server did not receive the right Authorization header from the KaggleDatasets client.")
92+
93+
def test_no_token_fails(self):
94+
env = EnvironmentVarGuard()
95+
env.unset(_KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME)
96+
with env:
97+
with self.assertRaises(CredentialError):
98+
client = KaggleDatasets()
99+
100+
def test_get_gcs_path_tpu_succeeds(self):
101+
def call_get_gcs_path():
102+
client = KaggleDatasets()
103+
gcs_path = client.get_gcs_path()
104+
self.assertEqual(gcs_path, _TPU_GCS_BUCKET)
105+
self._test_client(call_get_gcs_path,
106+
'/requests/CopyDatasetVersionToKnownGcsBucketRequest',
107+
{'MountSlug': None, 'IntegrationType': 2, 'JWE': _TEST_JWT},
108+
is_tpu=True)
109+
110+
def test_get_gcs_path_automl_succeeds(self):
111+
def call_get_gcs_path():
112+
client = KaggleDatasets()
113+
gcs_path = client.get_gcs_path()
114+
self.assertEqual(gcs_path, _AUTOML_GCS_BUCKET)
115+
self._test_client(call_get_gcs_path,
116+
'/requests/CopyDatasetVersionToKnownGcsBucketRequest',
117+
{'MountSlug': None, 'IntegrationType': 1, 'JWE': _TEST_JWT},
118+
is_tpu=False)
119+
120+
def test_get_gcs_path_handles_unsuccessful(self):
121+
def call_get_gcs_path():
122+
client = KaggleDatasets()
123+
with self.assertRaises(BackendError):
124+
gcs_path = client.get_gcs_path()
125+
self._test_client(call_get_gcs_path,
126+
'/requests/CopyDatasetVersionToKnownGcsBucketRequest',
127+
{'MountSlug': None, 'IntegrationType': 2, 'JWE': _TEST_JWT},
128+
is_tpu=True,
129+
success=False)

0 commit comments

Comments
 (0)
Back | FazBrowse Home | New Git URL