import argparse
import requests
import time

from polyaxon_client.client import PolyaxonClient
from polyaxon_client.exceptions import PolyaxonHTTPError


def create_experiment_auth_context(experiment_name: str):
    parts = experiment_name.split('.')
    client = PolyaxonClient()
    client.auth.login_experiment_impersonate_token(username=parts[0],
                                                   project_name=parts[1],
                                                   experiment_id=parts[-1],
                                                   internal_token=client.api_config.token)


def create_job_auth_context(job_name: str):
    parts = job_name.split('.')
    client = PolyaxonClient()
    client.auth.login_job_impersonate_token(username=parts[0],
                                            project_name=parts[1],
                                            job_id=parts[-1],
                                            internal_token=client.api_config.token)


def create_notebook_auth_context(job_name: str):
    parts = job_name.split('.')
    client = PolyaxonClient()
    client.auth.login_notebook_impersonate_token(username=parts[0],
                                                 project_name=parts[1],
                                                 internal_token=client.api_config.token)


def _create_auth_context():
    if entity == 'experiment':
        create_experiment_auth_context(experiment_name=entity_name)
    elif entity == 'job':
        create_job_auth_context(job_name=entity_name)
    elif entity == 'notebook':
        create_notebook_auth_context(job_name=entity_name)
    else:
        raise ValueError('Entity `{}` not recognized'.format(entity))

    return True


def create_auth_context():
    retry = 0
    done = False
    while not done and retry < 3:
        try:
            done = _create_auth_context()
        except (requests.exceptions.ConnectionError,
                requests.exceptions.RequestException,
                requests.exceptions.Timeout,
                requests.exceptions.HTTPError,
                PolyaxonHTTPError):
            retry += 1
            print('Could not establish connection, retrying ...')
            time.sleep(sleep_interval)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument(
        '--entity',
        type=str
    )
    parser.add_argument(
        '--entity_name',
        type=str
    )
    parser.add_argument(
        '--sleep_interval',
        default=2,
        type=int
    )
    parser.set_defaults(nocache=False)
    args = parser.parse_args()
    arguments = args.__dict__

    entity = arguments.pop('entity')
    entity_name = arguments.pop('entity_name')
    sleep_interval = arguments.pop('sleep_interval')

    create_auth_context()

@@ -38,16 +38,42 @@
 t time%0A%0A
+from pathlib import Path%0A%0A
 from pol
@@ -171,16 +171,77 @@
 TPError%0A
+from polyaxon_client.settings import CONTEXT_AUTH_TOKEN_PATH%0A
 %0A%0Adef cr
@@ -1424,24 +1424,91 @@
 _context():%0A
+    if Path(CONTEXT_AUTH_TOKEN_PATH).exists():%0A        return True%0A
     if entit

Add eventual pod restart handling for init process