def get_repo_context_args(
name: str,
url: str,
revision: str,
mount_path: str,
connection: str = None,
flags: List[str] = None,
) -> List[str]:
if not name:
raise PolypodException("A repo name is required to create a repo context.")
if not url:
raise PolypodException("A repo url is required to create a repo context.")
args = ["--repo-path={}/{}".format(mount_path, name), "--url={}".format(url)]
if revision:
args.append("--revision={}".format(revision))
if connection:
args.append("--connection={}".format(connection))
flags = to_list(flags, check_none=True)
if flags:
args.append("--flags={}".format(ujson.dumps(flags)))
return args
def get_str_var(value: Any) -> str:
if value is not None and not isinstance(value, str):
try:
value = json.dumps(value)
except (ValueError, TypeError) as e:
raise PolypodException(e)
return value or ""
def get_artifacts_path_container(
polyaxon_init: V1PolyaxonInitContainer,
artifacts_store: V1ConnectionType,
run_path: str,
auto_resume: bool,
) -> Optional[k8s_schemas.V1Container]:
if not artifacts_store:
raise PolypodException("Init artifacts container requires a store.")
init_args = init_artifact_context_args(run_path=run_path)
if auto_resume:
init_args.append(
get_volume_args(
store=artifacts_store,
mount_path=CONTEXT_MOUNT_ARTIFACTS,
artifacts=V1ArtifactsType(dirs=[run_path]),
))
container_name = generate_container_name(INIT_ARTIFACTS_CONTAINER_PREFIX,
DEFAULT)
container = k8s_schemas.V1Container(name=container_name)
return get_base_store_container(
container_name=container_name,
container=container,
polyaxon_init=polyaxon_init,
store=artifacts_store,
env=[],
env_from=[],
volume_mounts=[get_artifacts_context_mount()],
# If we are dealing with a volume we need to make sure the path exists for the user
# We also clean the path if this is not a resume run
args=[" ".join(init_args)],
)
def get_artifacts_path_container(
polyaxon_init: V1PolyaxonInitContainer,
artifacts_store: V1ConnectionType,
run_path: str,
clean: bool = True,
) -> Optional[k8s_schemas.V1Container]:
if not artifacts_store:
raise PolypodException("Init artifacts container requires a store.")
init_args = init_artifact_context_args(run_path=run_path)
if not artifacts_store.is_bucket:
artifacts_path = get_path(artifacts_store.store_path, run_path)
init_args.append(
get_artifacts_store_args(artifacts_path=artifacts_path,
clean=clean))
container_name = INIT_ARTIFACTS_CONTAINER.format(DEFAULT)
container = k8s_schemas.V1Container(name=container_name)
return get_base_store_container(
container_name=container_name,
container=container,
polyaxon_init=polyaxon_init,
store=artifacts_store,
env=[],
env_from=[],
volume_mounts=[get_artifacts_context_mount()],
# If we are dealing with a volume we need to make sure the path exists for the user
# We also clean the path if this is not a resume run
args=[" ".join(init_args)],
is_artifact_store=True,
)
def get_service_env_vars(
header: str,
service_header: str,
include_secret_key: bool,
include_internal_token: bool,
include_agent_token: bool,
authentication_type: str,
polyaxon_default_secret_ref: str,
polyaxon_agent_secret_ref: str,
api_host: str,
api_version: str,
run_instance: str,
) -> List[k8s_schemas.V1EnvVar]:
env_vars = get_base_env_vars() + [
get_env_var(name=POLYAXON_KEYS_HOST, value=api_host),
get_env_var(name=POLYAXON_KEYS_IS_MANAGED, value=True),
get_env_var(name=POLYAXON_KEYS_API_VERSION, value=api_version),
get_run_instance_env_var(run_instance),
]
if header:
env_vars.append(
get_env_var(
name=POLYAXON_KEYS_HEADER,
value=PolyaxonServiceHeaders.get_header(header),
))
if service_header:
env_vars.append(
get_env_var(name=POLYAXON_KEYS_HEADER_SERVICE,
value=service_header))
if include_secret_key:
env_vars.append(
get_from_secret(
key_name=POLYAXON_KEYS_SECRET_KEY,
secret_key_name=POLYAXON_KEYS_SECRET_KEY,
secret_ref_name=polyaxon_default_secret_ref,
))
internal = False
if include_internal_token and polyaxon_default_secret_ref:
internal = True
env_vars.append(
get_from_secret(
POLYAXON_KEYS_SECRET_INTERNAL_TOKEN,
POLYAXON_KEYS_SECRET_INTERNAL_TOKEN,
secret_ref_name=polyaxon_default_secret_ref,
))
if include_agent_token and polyaxon_agent_secret_ref:
if internal:
raise PolypodException(
"A service cannot have internal token and agent token.")
env_vars.append(
get_from_secret(
POLYAXON_KEYS_AUTH_TOKEN,
POLYAXON_KEYS_AUTH_TOKEN,
secret_ref_name=polyaxon_agent_secret_ref,
))
if authentication_type:
env_vars.append(
get_env_var(name=POLYAXON_KEYS_AUTHENTICATION_TYPE,
value=authentication_type))
return env_vars
def get_base_store_container(
container: Optional[k8s_schemas.V1Container],
container_name: str,
polyaxon_init: V1PolyaxonInitContainer,
store: V1ConnectionType,
env: List[k8s_schemas.V1EnvVar],
env_from: List[k8s_schemas.V1EnvFromSource],
volume_mounts: List[k8s_schemas.V1VolumeMount],
args: List[str],
is_artifact_store: Optional[bool] = False,
) -> Optional[k8s_schemas.V1Container]:
env = env or []
env_from = env_from or []
volume_mounts = volume_mounts or []
# Artifact store needs to allow init the contexts as well, so the store is not required
if not is_artifact_store and not store:
raise PolypodException("Init store container requires a store")
secret = None
if store.is_bucket:
if not is_artifact_store:
secret = store.get_secret()
volume_mounts = volume_mounts + to_list(
get_mount_from_resource(resource=secret), check_none=True)
env = env + to_list(get_items_from_secret(secret=secret),
check_none=True)
env_from = env_from + to_list(get_env_from_secret(secret=secret),
check_none=True)
env += to_list(get_connection_env_var(connection=store,
secret=secret),
check_none=True)
config_map = store.get_config_map()
volume_mounts = volume_mounts + to_list(
get_mount_from_resource(resource=config_map), check_none=True)
env = env + to_list(
get_items_from_config_map(config_map=config_map),
check_none=True)
env_from = env_from + to_list(
get_env_from_config_map(config_map=config_map),
check_none=True)
else:
volume_mounts = volume_mounts + to_list(
get_mount_from_store(store=store), check_none=True)
env += to_list(get_connection_env_var(connection=store, secret=secret),
check_none=True)
return patch_container(
container=container,
name=container_name,
image=polyaxon_init.get_image(),
image_pull_policy=polyaxon_init.image_pull_policy,
command=["/bin/sh", "-c"],
args=args,
env=env,
env_from=env_from,
resources=polyaxon_init.get_resources(),
volume_mounts=volume_mounts,
)
def get_env_var(name: str, value: Any) -> k8s_schemas.V1EnvVar:
if not isinstance(value, str):
try:
value = json.dumps(value)
except (ValueError, TypeError) as e:
raise PolypodException(e)
return k8s_schemas.V1EnvVar(name=name, value=value)
def is_valid(self):
if not self.SPEC_KIND:
raise PolypodException(
"Please make sure that a spawner subclass has a valid SPEC_KIND"
)
if not self.GROUP:
raise PolypodException(
"Please make sure that a spawner subclass has a valid GROUP")
if not self.API_VERSION:
raise PolypodException(
"Please make sure that a spawner subclass has a valid API_VERSION"
)
if not self.PLURAL:
raise PolypodException(
"Please make sure that a spawner subclass has a valid PLURAL")
if not self.K8S_ANNOTATIONS_KIND:
raise PolypodException(
"Please make sure that a spawner subclass has a valid K8S_ANNOTATIONS_KIND"
)
if not self.K8S_LABELS_COMPONENT:
raise PolypodException(
"Please make sure that a spawner subclass has a valid K8S_LABELS_COMPONENT"
)
if not self.K8S_LABELS_PART_OF:
raise PolypodException(
"Please make sure that a spawner subclass has a valid K8S_LABELS_PART_OF"
)
def get_pod_spec(
resource_name: str,
namespace: str,
main_container: k8s_schemas.V1Container,
sidecar_containers: Optional[List[k8s_schemas.V1Container]],
init_containers: Optional[List[k8s_schemas.V1Container]],
environment: Optional[V1Environment],
labels: Dict[str, str],
volumes: Optional[List[k8s_schemas.V1Volume]],
) -> Tuple[k8s_schemas.V1ObjectMeta, k8s_schemas.V1PodSpec]:
if not main_container:
raise PolypodException("A main container is required")
environment = environment or V1Environment()
metadata = k8s_schemas.V1ObjectMeta(
name=resource_name,
namespace=namespace,
labels=labels,
annotations=environment.annotations,
)
init_containers = to_list(init_containers, check_none=True)
containers = [main_container] + to_list(sidecar_containers,
check_none=True)
image_pull_secrets = None
if environment.image_pull_secrets:
image_pull_secrets = [
k8s_schemas.V1LocalObjectReference(name=i)
for i in environment.image_pull_secrets
]
pod_spec = k8s_schemas.V1PodSpec(
init_containers=init_containers,
containers=containers,
volumes=volumes,
restart_policy=environment.restart_policy,
image_pull_secrets=image_pull_secrets,
security_context=environment.security_context,
service_account_name=environment.service_account_name,
node_selector=environment.node_selector,
tolerations=environment.tolerations,
affinity=environment.affinity,
dns_config=environment.dns_config,
dns_policy=environment.dns_policy,
host_aliases=environment.host_aliases,
host_network=environment.host_network,
host_pid=environment.host_pid,
node_name=environment.node_name,
priority=environment.priority,
priority_class_name=environment.priority_class_name,
scheduler_name=environment.scheduler_name,
)
return metadata, pod_spec
def get_kv_env_vars(kv_env_vars: List[List]) -> List[k8s_schemas.V1EnvVar]:
env_vars = []
if not kv_env_vars:
return env_vars
for kv_env_var in kv_env_vars:
if not kv_env_var or not len(kv_env_var) == 2:
raise PolypodException(
"Received a wrong a key value env var `{}`".format(kv_env_var))
env_vars.append(get_env_var(name=kv_env_var[0], value=kv_env_var[1]))
return env_vars
def get_repo_context_args(name: str,
url: str,
revision: str,
mount_path: str,
connection: str = None) -> List[str]:
if not name:
raise PolypodException(
"A repo name is required to create a repo context.")
if not url:
raise PolypodException(
"A repo url is required to create a repo context.")
args = [
"--repo_path={}/{}".format(mount_path, name), "--url={}".format(url)
]
if revision:
args.append("--revision={}".format(revision))
if connection:
args.append("--connection={}".format(connection))
return args
def get_custom_init_container(
connection: V1ConnectionType,
contexts: PluginsContextsSpec,
container: Optional[k8s_schemas.V1Container],
env: List[k8s_schemas.V1EnvVar] = None,
mount_path: str = None,
) -> k8s_schemas.V1Container:
if not connection:
raise PolypodException(
"A connection is required to create a repo context.")
volume_name = (get_volume_name(mount_path)
if mount_path else constants.CONTEXT_VOLUME_ARTIFACTS)
mount_path = mount_path or CONTEXT_MOUNT_ARTIFACTS
volume_mounts = [
get_connections_context_mount(name=volume_name, mount_path=mount_path)
]
if contexts and contexts.auth:
volume_mounts.append(get_auth_context_mount(read_only=True))
env = to_list(env, check_none=True)
env_from = []
secret = connection.get_secret()
if secret:
volume_mounts += to_list(get_mount_from_resource(resource=secret),
check_none=True)
env += to_list(get_items_from_secret(secret=secret), check_none=True)
env_from = to_list(get_env_from_secret(secret=secret), check_none=True)
env += to_list(get_connection_env_var(connection=connection,
secret=secret),
check_none=True)
config_map = connection.get_config_map()
if config_map:
volume_mounts += to_list(get_mount_from_resource(resource=config_map),
check_none=True)
env += to_list(get_items_from_config_map(config_map=config_map),
check_none=True)
env_from = to_list(get_env_from_config_map(config_map=config_map),
check_none=True)
container_name = container.name or generate_container_name(
INIT_CUSTOM_CONTAINER_PREFIX, connection.name)
return patch_container(
container=container,
name=container_name,
env=env,
env_from=env_from,
volume_mounts=volume_mounts,
)
def get_env_vars(
contexts: PluginsContextsSpec,
log_level: str,
kv_env_vars: List[List],
artifacts_store_name: str,
connections: Iterable[V1ConnectionType],
secrets: Iterable[V1K8sResourceType],
config_maps: Iterable[V1K8sResourceType],
) -> List[k8s_schemas.V1EnvVar]:
env_vars = []
connections = connections or []
if log_level:
env_vars.append(get_env_var(name=POLYAXON_KEYS_LOG_LEVEL, value=log_level))
if contexts and contexts.collect_artifacts:
env_vars.append(get_env_var(name=POLYAXON_KEYS_COLLECT_ARTIFACTS, value=True))
if contexts and contexts.collect_resources:
env_vars.append(get_env_var(name=POLYAXON_KEYS_COLLECT_RESOURCES, value=True))
if artifacts_store_name:
env_vars.append(
get_env_var(
name=POLYAXON_KEYS_ARTIFACTS_STORE_NAME, value=artifacts_store_name
)
)
# Add connection env vars information
for connection in connections:
try:
secret = connection.get_secret()
env_vars += to_list(
get_connection_env_var(connection=connection, secret=secret),
check_none=True,
)
except PolyaxonSchemaError as e:
raise PolypodException("Error resolving secrets: %s" % e) from e
env_vars += get_kv_env_vars(kv_env_vars)
env_vars += get_env_vars_from_k8s_resources(
secrets=secrets, config_maps=config_maps
)
return env_vars
def get_main_container(
container_id: str,
main_container: k8s_schemas.V1Container,
volume_mounts: List[k8s_schemas.V1VolumeMount],
contexts: PluginsContextsSpec,
artifacts_store: Optional[V1ConnectionType],
init: Optional[List[V1Init]],
connections: Optional[List[str]],
connection_by_names: Dict[str, V1ConnectionType],
secrets: Optional[Iterable[V1K8sResourceType]],
config_maps: Optional[Iterable[V1K8sResourceType]],
log_level: str,
run_path: Optional[str],
kv_env_vars=None,
env: List[k8s_schemas.V1EnvVar] = None,
ports: List[int] = None,
) -> k8s_schemas.V1Container:
"""Pod job container for task."""
connections = connections or []
connection_by_names = connection_by_names or {}
secrets = secrets or []
config_maps = config_maps or []
if artifacts_store and not run_path:
raise PolypodException("Run path is required for main container.")
if artifacts_store and not contexts.collect_artifacts:
if artifacts_store.name not in connection_by_names:
connection_by_names[artifacts_store.name] = artifacts_store
connections.append(artifacts_store.name)
requested_connections = [connection_by_names[c] for c in connections]
requested_config_maps = get_requested_config_maps(
config_maps=config_maps, connections=requested_connections)
requested_secrets = get_requested_secrets(
secrets=secrets, connections=requested_connections)
# Mounts
volume_mounts = to_list(volume_mounts, check_none=True)
volume_mounts = volume_mounts + get_volume_mounts(
contexts=contexts,
init=init,
connections=requested_connections,
secrets=requested_secrets,
config_maps=requested_config_maps,
)
# Env vars
env = to_list(env, check_none=True)
env = env + get_env_vars(
contexts=contexts,
log_level=log_level,
kv_env_vars=kv_env_vars,
connections=requested_connections,
secrets=requested_secrets,
config_maps=requested_config_maps,
)
# Env from
env_from = get_env_from_k8s_resources(secrets=requested_secrets,
config_maps=requested_config_maps)
ports = [
k8s_schemas.V1ContainerPort(container_port=port)
for port in to_list(ports, check_none=True)
]
return patch_container(
container=main_container,
name=container_id,
env=env,
env_from=env_from,
volume_mounts=volume_mounts,
ports=ports or None,
)
def get_service_env_vars(
header: str,
service_header: str,
include_secret_key: bool,
include_internal_token: bool,
include_agent_token: bool,
authentication_type: str,
polyaxon_default_secret_ref: str,
polyaxon_agent_secret_ref: str,
api_host: str,
api_version: str,
run_instance: str,
) -> List[k8s_schemas.V1EnvVar]:
env_vars = [
get_env_var(name=POLYAXON_KEYS_HOST, value=api_host),
get_env_var(
name=POLYAXON_KEYS_API_HOST, value=api_host
), # TODO: Remove in v1.2
get_env_var(name=POLYAXON_KEYS_IS_MANAGED, value=True),
get_env_var(name=POLYAXON_KEYS_API_VERSION, value=api_version),
# Pod info
get_from_field_ref(
name=POLYAXON_KEYS_K8S_NODE_NAME, field_path="spec.nodeName"
),
get_from_field_ref(
name=POLYAXON_KEYS_K8S_NAMESPACE, field_path="metadata.namespace"
),
get_from_field_ref(name=POLYAXON_KEYS_K8S_POD_ID, field_path="metadata.name"),
get_run_instance_env_var(run_instance),
]
if header:
env_vars.append(
get_env_var(
name=POLYAXON_KEYS_HEADER,
value=PolyaxonServiceHeaders.get_header(header),
)
)
if service_header:
env_vars.append(
get_env_var(name=POLYAXON_KEYS_HEADER_SERVICE, value=service_header)
)
if include_secret_key:
env_vars.append(
get_from_secret(
key_name=POLYAXON_KEYS_SECRET_KEY,
secret_key_name=POLYAXON_KEYS_SECRET_KEY,
secret_ref_name=polyaxon_default_secret_ref,
)
)
internal = False
if include_internal_token and polyaxon_default_secret_ref:
internal = True
env_vars.append(
get_from_secret(
POLYAXON_KEYS_SECRET_INTERNAL_TOKEN,
POLYAXON_KEYS_SECRET_INTERNAL_TOKEN,
secret_ref_name=polyaxon_default_secret_ref,
)
)
if include_agent_token and polyaxon_agent_secret_ref:
if internal:
raise PolypodException(
"A service cannot have internal token and agent token."
)
env_vars.append(
get_from_secret(
POLYAXON_KEYS_AUTH_TOKEN,
POLYAXON_KEYS_AUTH_TOKEN,
secret_ref_name=polyaxon_agent_secret_ref,
)
)
if authentication_type:
env_vars.append(
get_env_var(
name=POLYAXON_KEYS_AUTHENTICATION_TYPE, value=authentication_type
)
)
return env_vars
def get_sidecar_container(
container_id: str,
polyaxon_sidecar: V1PolyaxonSidecarContainer,
env: List[k8s_schemas.V1EnvVar],
artifacts_store: V1ConnectionType,
contexts: PluginsContextsSpec,
run_path: Optional[str],
) -> Optional[k8s_schemas.V1Container]:
if artifacts_store and not contexts:
raise PolypodException(
"Logs/artifacts store was passed and contexts was not passed.")
has_artifacts = artifacts_store and contexts.collect_artifacts
has_logs = artifacts_store and contexts.collect_logs
if not has_logs and not has_artifacts:
# No sidecar
return None
if (has_artifacts or has_logs) and not run_path:
raise PolypodException(
"Logs store / outputs store must have a run_path.")
env = get_sidecar_env_vars(
env_vars=env,
container_id=container_id,
artifacts_store_name=artifacts_store.name,
)
volume_mounts = get_mounts(
use_auth_context=contexts.auth,
use_artifacts_context=has_artifacts,
use_docker_context=False,
use_shm_context=False,
)
sidecar_args = get_sidecar_args(
container_id=container_id,
sleep_interval=polyaxon_sidecar.sleep_interval,
sync_interval=polyaxon_sidecar.sync_interval,
)
env_from = []
secret = None
if artifacts_store.is_bucket:
secret = artifacts_store.get_secret()
volume_mounts += to_list(get_mount_from_resource(resource=secret),
check_none=True)
env += to_list(get_items_from_secret(secret=secret), check_none=True)
env_from += to_list(get_env_from_secret(secret=secret),
check_none=True)
config_map = artifacts_store.get_config_map()
volume_mounts += to_list(get_mount_from_resource(resource=config_map),
check_none=True)
env += to_list(get_items_from_config_map(config_map=config_map),
check_none=True)
env_from += to_list(get_env_from_config_map(config_map=config_map),
check_none=True)
else:
volume_mounts += to_list(get_mount_from_store(store=artifacts_store),
check_none=True)
env += to_list(
get_connection_env_var(connection=artifacts_store, secret=secret),
check_none=True,
)
return k8s_schemas.V1Container(
name=SIDECAR_CONTAINER,
image=polyaxon_sidecar.get_image(),
image_pull_policy=polyaxon_sidecar.image_pull_policy,
command=["/bin/bash", "-c"],
args=[sidecar_args],
env=env,
env_from=env_from,
resources=polyaxon_sidecar.get_resources(),
volume_mounts=volume_mounts,
)
def get_git_init_container(
polyaxon_init: V1PolyaxonInitContainer,
connection: V1ConnectionType,
contexts: PluginsContextsSpec,
container: Optional[k8s_schemas.V1Container] = None,
env: List[k8s_schemas.V1EnvVar] = None,
mount_path: str = None,
track: bool = False,
) -> k8s_schemas.V1Container:
if not connection:
raise PolypodException(
"A connection is required to create a repo context.")
if not container:
container = k8s_schemas.V1Container(name=generate_container_name(
INIT_GIT_CONTAINER_PREFIX, connection.name), )
volume_name = (get_volume_name(mount_path)
if mount_path else constants.CONTEXT_VOLUME_ARTIFACTS)
mount_path = mount_path or CONTEXT_MOUNT_ARTIFACTS
volume_mounts = [
get_connections_context_mount(name=volume_name, mount_path=mount_path)
]
if contexts and contexts.auth:
volume_mounts.append(get_auth_context_mount(read_only=True))
env = to_list(env, check_none=True)
env_from = []
secret = connection.get_secret()
if secret:
volume_mounts += to_list(get_mount_from_resource(resource=secret),
check_none=True)
env += to_list(get_items_from_secret(secret=secret), check_none=True)
env_from = to_list(get_env_from_secret(secret=secret), check_none=True)
env += to_list(get_connection_env_var(connection=connection,
secret=secret),
check_none=True)
config_map = connection.get_config_map()
if config_map:
volume_mounts += to_list(get_mount_from_resource(resource=config_map),
check_none=True)
env += to_list(get_items_from_config_map(config_map=config_map),
check_none=True)
env_from = to_list(get_env_from_config_map(config_map=config_map),
check_none=True)
args = get_repo_context_args(
name=connection.name,
url=connection.schema.url,
revision=connection.schema.revision,
mount_path=mount_path,
connection=connection.name if track else None,
)
return patch_container(
container=container,
name=generate_container_name(INIT_GIT_CONTAINER_PREFIX,
connection.name),
image=polyaxon_init.get_image(),
image_pull_policy=polyaxon_init.image_pull_policy,
command=["polyaxon", "initializer", "git"],
args=args,
env=env,
env_from=env_from,
volume_mounts=volume_mounts,
resources=polyaxon_init.get_resources(),
)