def test_get_mount_from_store(self):
# Non bucket stores
assert get_mount_from_store(store=None) is None
store = V1ConnectionType(
name="test",
kind=V1ConnectionKind.S3,
schema=V1BucketConnection(bucket="s3//:foo"),
)
assert get_mount_from_store(store=store) is None
# Claim store
store = V1ConnectionType(
name="test",
kind=V1ConnectionKind.VOLUME_CLAIM,
schema=V1ClaimConnection(mount_path="/tmp",
volume_claim="test",
read_only=True),
)
mount = get_mount_from_store(store=store)
assert mount.name == store.name
assert mount.mount_path == store.schema.mount_path
assert mount.read_only == store.schema.read_only
# Host path
store = V1ConnectionType(
name="test",
kind=V1ConnectionKind.HOST_PATH,
schema=V1HostPathConnection(mount_path="/tmp",
host_path="/tmp",
read_only=True),
)
mount = get_mount_from_store(store=store)
assert mount.name == store.name
assert mount.mount_path == store.schema.mount_path
assert mount.read_only == store.schema.read_only
def test_single_connections(self):
self.assert_single_store(store=self.s3_store, results=[])
self.assert_single_store(store=self.gcs_store, results=[])
self.assert_single_store(store=self.az_store, results=[])
self.assert_single_store(
store=self.claim_store,
results=[get_mount_from_store(store=self.claim_store)],
)
self.assert_single_store(
store=self.host_path_store,
results=[get_mount_from_store(store=self.host_path_store)],
)
# Managed versions
volume_name = get_volume_name("/test")
self.assert_single_init_store(
store=self.s3_store,
results=[
get_connections_context_mount(name=volume_name,
mount_path="/test")
],
)
self.assert_single_init_store(
store=self.gcs_store,
results=[
get_connections_context_mount(name=volume_name,
mount_path="/test")
],
)
self.assert_single_init_store(
store=self.az_store,
results=[
get_connections_context_mount(name=volume_name,
mount_path="/test")
],
)
self.assert_single_init_store(
store=self.claim_store,
results=[
get_connections_context_mount(name=volume_name,
mount_path="/test")
],
)
self.assert_single_init_store(
store=self.host_path_store,
results=[
get_connections_context_mount(name=volume_name,
mount_path="/test")
],
)
def test_get_base_container(self):
store = V1ConnectionType(
name="test_claim",
kind=V1ConnectionKind.VOLUME_CLAIM,
schema=V1ClaimConnection(mount_path="/tmp",
volume_claim="test",
read_only=True),
)
env = [get_env_var(name="key", value="value")]
env_from = [k8s_schemas.V1EnvFromSource(secret_ref={"name": "ref"})]
mounts = [k8s_schemas.V1VolumeMount(name="test", mount_path="/test")]
container = get_base_store_container(
container=k8s_schemas.V1Container(name="init"),
container_name="init",
polyaxon_init=V1PolyaxonInitContainer(
image="foo/foo",
image_tag="",
image_pull_policy="IfNotPresent"),
store=store,
env=env,
env_from=env_from,
volume_mounts=mounts,
args=["test"],
)
assert container.name == "init"
assert container.image == "foo/foo"
assert container.image_pull_policy == "IfNotPresent"
assert container.command == ["/bin/sh", "-c"]
assert container.args == ["test"]
assert container.env == env
assert container.env_from == env_from
assert container.resources is not None
assert container.volume_mounts == mounts + [
get_mount_from_store(store=store)
]
def test_get_base_store_container_with_mount_store(self):
claim_store = V1ConnectionType(
name="test_claim",
kind=V1ConnectionKind.VOLUME_CLAIM,
schema=V1ClaimConnection(mount_path="/tmp",
volume_claim="test",
read_only=True),
)
container = get_base_store_container(
container=k8s_schemas.V1Container(name="init"),
container_name="init",
polyaxon_init=V1PolyaxonInitContainer(image_tag=""),
store=claim_store,
env=None,
env_from=None,
volume_mounts=None,
args=None,
)
assert container.name == "init"
assert container.image == "polyaxon/polyaxon-init"
assert container.image_pull_policy is None
assert container.command == ["/bin/sh", "-c"]
assert container.args is None
assert container.env == get_connection_env_var(connection=claim_store,
secret=None)
assert container.env_from == []
assert container.resources is not None
assert container.volume_mounts == [
get_mount_from_store(store=claim_store)
]
def test_get_store_container_mount_stores(self):
# Managed store
store = V1ConnectionType(
name="test_claim",
kind=V1ConnectionKind.VOLUME_CLAIM,
schema=V1ClaimConnection(
mount_path="/tmp", volume_claim="test", read_only=True
),
)
container = get_store_container(
polyaxon_init=V1PolyaxonInitContainer(
image="foo/foo", image_tag="foo", image_pull_policy="IfNotPresent"
),
connection=store,
artifacts=None,
)
mount_path = CONTEXT_MOUNT_ARTIFACTS_FORMAT.format(store.name)
assert container.name == INIT_ARTIFACTS_CONTAINER.format(store.name)
assert container.image == "foo/foo:foo"
assert container.image_pull_policy == "IfNotPresent"
assert container.command == ["/bin/sh", "-c"]
assert container.args == [
get_volume_args(store=store, mount_path=mount_path, artifacts=None)
]
assert container.env == get_connection_env_var(connection=store, secret=None)
assert container.env_from == []
assert container.resources is not None
assert container.volume_mounts == [
get_connections_context_mount(
name=constants.CONTEXT_VOLUME_ARTIFACTS,
mount_path=CONTEXT_MOUNT_ARTIFACTS_FORMAT.format(store.name),
),
get_mount_from_store(store=store),
]
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_volume_mounts(
contexts: PluginsContextsSpec,
init: Optional[List[V1Init]],
connections: Iterable[V1ConnectionType],
secrets: Iterable[V1K8sResourceType],
config_maps: Iterable[V1K8sResourceType] = None,
) -> List[k8s_schemas.V1VolumeMount]:
init = init or []
connections = connections or []
secrets = secrets or []
config_maps = config_maps or []
volume_mounts = []
volume_names = set()
if contexts and contexts.collect_artifacts:
volume_mounts += to_list(
get_artifacts_context_mount(read_only=False), check_none=True
)
volume_names.add(constants.CONTEXT_VOLUME_ARTIFACTS)
for init_connection in init:
volume_name = (
get_volume_name(init_connection.path)
if init_connection.path
else constants.CONTEXT_VOLUME_ARTIFACTS
)
mount_path = init_connection.path or CONTEXT_MOUNT_ARTIFACTS
if volume_name in volume_names:
continue
volume_names.add(volume_name)
volume_mounts += to_list(
get_connections_context_mount(name=volume_name, mount_path=mount_path),
check_none=True,
)
for store in connections:
volume_mounts += to_list(get_mount_from_store(store=store), check_none=True)
for secret in secrets:
volume_mounts += to_list(
get_mount_from_resource(resource=secret), check_none=True
)
for config_map in config_maps:
volume_mounts += to_list(
get_mount_from_resource(resource=config_map), check_none=True
)
return volume_mounts
def test_get_sidecar_container_host_paths(self):
artifacts_store = V1ConnectionType(
name="plx-outputs",
kind=V1ConnectionKind.HOST_PATH,
schema=V1HostPathConnection(mount_path="/tmp/plx/outputs",
host_path="/tmp/plx/outputs"),
)
contexts = PluginsContextsSpec(
auth=True,
docker=False,
shm=False,
collect_logs=True,
collect_artifacts=True,
collect_resources=True,
auto_resume=True,
sync_statuses=True,
external_host=False,
sidecar=None,
)
sidecar = get_sidecar_container(
container_id=MAIN_JOB_CONTAINER,
polyaxon_sidecar=V1PolyaxonSidecarContainer(
image="foo",
image_pull_policy="sdf",
sleep_interval=2,
sync_interval=212,
),
env=[],
artifacts_store=artifacts_store,
contexts=contexts,
run_path="test",
)
assert sidecar.volume_mounts == [
get_auth_context_mount(read_only=True),
get_artifacts_context_mount(read_only=False),
get_mount_from_store(store=artifacts_store),
]
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 test_get_sidecar_container_with_managed_mount_outputs_logs_store(self):
env_vars = [
get_env_var(name="key1", value="value1"),
get_env_var(name="key2", value="value2"),
]
mount_managed_store = V1ConnectionType(
name="test_path",
kind=V1ConnectionKind.HOST_PATH,
schema=V1HostPathConnection(mount_path="/tmp", host_path="/tmp"),
secret=None,
)
# logs and outputs
sidecar = get_sidecar_container(
container_id=MAIN_JOB_CONTAINER,
env=env_vars,
polyaxon_sidecar=V1PolyaxonSidecarContainer(
image="sidecar/sidecar",
image_tag="",
image_pull_policy="IfNotPresent",
sleep_interval=213,
sync_interval=212,
),
artifacts_store=mount_managed_store,
contexts=PluginsContextsSpec.from_config(None, default_auth=True),
run_path="test",
)
assert sidecar.name == SIDECAR_CONTAINER
assert sidecar.image == "sidecar/sidecar"
assert sidecar.image_pull_policy == "IfNotPresent"
assert sidecar.command == ["polyaxon", "sidecar"]
assert sidecar.args == get_sidecar_args(
container_id=MAIN_JOB_CONTAINER,
sleep_interval=213,
sync_interval=212,
monitor_logs=False,
)
assert sidecar.env == get_sidecar_env_vars(
env_vars=env_vars,
container_id=MAIN_JOB_CONTAINER,
artifacts_store_name=mount_managed_store.name,
) + get_connection_env_var(connection=mount_managed_store, secret=None)
assert sidecar.env_from == []
assert sidecar.resources == get_sidecar_resources()
assert sidecar.volume_mounts == [
get_auth_context_mount(read_only=True),
get_artifacts_context_mount(read_only=False),
get_mount_from_store(store=mount_managed_store),
]
# logs and no outputs
sidecar = get_sidecar_container(
container_id=MAIN_JOB_CONTAINER,
env=env_vars,
polyaxon_sidecar=V1PolyaxonSidecarContainer(
image="sidecar/sidecar",
image_tag="",
image_pull_policy="IfNotPresent",
sleep_interval=213,
sync_interval=212,
),
artifacts_store=mount_managed_store,
contexts=PluginsContextsSpec.from_config(
V1Plugins(collect_artifacts=False, auth=True)),
run_path="test",
)
assert sidecar.name == SIDECAR_CONTAINER
assert sidecar.image == "sidecar/sidecar"
assert sidecar.image_pull_policy == "IfNotPresent"
assert sidecar.command == ["polyaxon", "sidecar"]
assert sidecar.args == get_sidecar_args(
container_id=MAIN_JOB_CONTAINER,
sleep_interval=213,
sync_interval=212,
monitor_logs=False,
)
assert sidecar.env == get_sidecar_env_vars(
env_vars=env_vars,
container_id=MAIN_JOB_CONTAINER,
artifacts_store_name=mount_managed_store.name,
) + get_connection_env_var(connection=mount_managed_store, secret=None)
assert sidecar.env_from == []
assert sidecar.resources == get_sidecar_resources()
assert sidecar.volume_mounts == [
get_auth_context_mount(read_only=True),
get_mount_from_store(store=mount_managed_store),
]
# outputs and no logs
sidecar = get_sidecar_container(
container_id=MAIN_JOB_CONTAINER,
env=env_vars,
polyaxon_sidecar=V1PolyaxonSidecarContainer(
image="sidecar/sidecar",
image_tag="",
image_pull_policy="IfNotPresent",
sleep_interval=213,
sync_interval=212,
),
artifacts_store=mount_managed_store,
contexts=PluginsContextsSpec.from_config(
V1Plugins(collect_logs=False, auth=True)),
run_path="test",
)
assert sidecar.name == SIDECAR_CONTAINER
assert sidecar.image == "sidecar/sidecar"
assert sidecar.image_pull_policy == "IfNotPresent"
assert sidecar.command == ["polyaxon", "sidecar"]
assert sidecar.args == get_sidecar_args(
container_id=MAIN_JOB_CONTAINER,
sleep_interval=213,
sync_interval=212,
monitor_logs=False,
)
assert sidecar.env == get_sidecar_env_vars(
env_vars=env_vars,
container_id=MAIN_JOB_CONTAINER,
artifacts_store_name=mount_managed_store.name,
) + get_connection_env_var(connection=mount_managed_store, secret=None)
assert sidecar.env_from == []
assert sidecar.resources == get_sidecar_resources()
assert sidecar.volume_mounts == [
get_auth_context_mount(read_only=True),
get_artifacts_context_mount(read_only=False),
get_mount_from_store(store=mount_managed_store),
]