diff --git a/sdk/python/kfp/_credentials.py b/sdk/python/kfp/_credentials.py index 2d545bf2321..f4e746600b6 100644 --- a/sdk/python/kfp/_credentials.py +++ b/sdk/python/kfp/_credentials.py @@ -12,14 +12,22 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os import abc __all__ = [ "TokenCredentials", + "ML_PIPELINE_SA_TOKEN_ENV", + "ML_PIPELINE_SA_TOKEN_PATH", + "ServiceAccountTokenVolumeCredentials", ] +ML_PIPELINE_SA_TOKEN_ENV = "ML_PIPELINE_SA_TOKEN_PATH" +ML_PIPELINE_SA_TOKEN_PATH = "/var/run/secrets/ml-pipeline/token" + + class TokenCredentials(object): __metaclass__ = abc.ABCMeta @@ -35,3 +43,36 @@ def refresh_api_key_hook(self, config): @abc.abstractmethod def get_token(self): raise NotImplementedError() + + +def read_sa_token(path=None): + """Read a ServiceAccount token found under some path.""" + token = None + with open(path, "r") as f: + token = f.read().strip() + return token + + +class ServiceAccountTokenVolumeCredentials(TokenCredentials): + """Audience-bound ServiceAccountToken in the local filesystem, that is + refreshed by the kubelet. + + The constructor of the class expects a filesystem path. + If not provided, it uses the path stored in the environment variable + defined in ML_PIPELINE_SA_TOKEN_ENV. + If the environment variable is also empty, it falls back to the path + specified in ML_PIPELINE_SA_TOKEN_PATH. + + This method of authentication is meant for use inside a Kubernetes cluster. + + Relevant documentation: + https://kubernetes.io/docs/tasks/configure-pod-container/configure-service-account/#service-account-token-volume-projection + """ + + def __init__(self, path=None): + self._token_path = (path + or os.getenv(ML_PIPELINE_SA_TOKEN_ENV) + or ML_PIPELINE_SA_TOKEN_PATH) + + def get_token(self): + return read_sa_token(self._token_path)