def test_get_init_containers_with_git_without_connection(self):
git1 = V1GitConnection(revision="test", url="https://test.com")
git2 = V1GitConnection(revision="test", url="https://test.com")
containers = self.converter.get_init_containers(
contexts=None,
artifacts_store=None,
init_connections=[
V1Init(git=git1,
container=k8s_schemas.V1Container(name="test")),
V1Init(git=git2, path="/test"),
],
init_containers=[],
connection_by_names={},
polyaxon_init=V1PolyaxonInitContainer(image="foo/foo"),
)
assert containers == [
get_git_init_container(
connection=V1ConnectionType(name=git1.get_name(),
kind=V1ConnectionKind.GIT,
schema=git1),
polyaxon_init=V1PolyaxonInitContainer(image="foo/foo"),
env=self.converter.get_init_service_env_vars(),
contexts=None,
),
get_git_init_container(
container=k8s_schemas.V1Container(name="test"),
connection=V1ConnectionType(name=git2.get_name(),
kind=V1ConnectionKind.GIT,
schema=git1),
mount_path="/test",
polyaxon_init=V1PolyaxonInitContainer(image="foo/foo"),
env=self.converter.get_init_service_env_vars(),
contexts=None,
),
]
def test_get_service_custom_resource(self):
main_container = k8s_schemas.V1Container(name="main")
sidecar_containers = [k8s_schemas.V1Container(name="sidecar")]
init_containers = [k8s_schemas.V1Container(name="init")]
termination = V1Termination(timeout=10)
environment = V1Environment(
labels={"foo": "bar"},
annotations={"foo": "bar"},
node_selector={"foo": "bar"},
node_name="foo",
restart_policy="never",
)
metadata, pod_spec = get_pod_spec(
namespace="default",
main_container=main_container,
sidecar_containers=sidecar_containers,
init_containers=init_containers,
resource_name="foo",
volumes=[],
environment=environment,
labels=environment.labels,
)
custom_object = {
"serviceSpec": {
"template":
get_pod_template_spec(metadata=metadata, pod_spec=pod_spec),
},
"termination": {
"activeDeadlineSeconds": termination.timeout
},
"collectLogs": True,
"syncStatuses": True,
"notifications": [],
}
expected_crd = get_custom_object(
namespace="default",
resource_name="foo",
kind="Operation",
api_version="core.polyaxon.com/v1",
labels={"foo": "bar"},
custom_object=custom_object,
)
crd = get_service_custom_resource(
namespace="default",
resource_name="foo",
main_container=main_container,
sidecar_containers=sidecar_containers,
init_containers=init_containers,
volumes=[],
termination=termination,
environment=environment,
labels=environment.labels,
collect_logs=True,
sync_statuses=True,
notifications=None,
ports=[],
)
assert crd == expected_crd
def get_replica(self, environment):
main_container = k8s_schemas.V1Container(name="main")
sidecar_containers = [k8s_schemas.V1Container(name="sidecar")]
init_containers = [k8s_schemas.V1Container(name="init")]
replica = ReplicaSpec(
volumes=[],
init_containers=init_containers,
sidecar_containers=sidecar_containers,
main_container=main_container,
labels=environment.labels,
environment=environment,
num_replicas=12,
)
metadata, pod_spec = get_pod_spec(
namespace="default",
main_container=main_container,
sidecar_containers=sidecar_containers,
init_containers=init_containers,
resource_name="foo",
volumes=[],
environment=environment,
labels=environment.labels,
)
replica_template = {
"replicas":
replica.num_replicas,
"restartPolicy":
pod_spec.restart_policy,
"template":
get_pod_template_spec(metadata=metadata, pod_spec=pod_spec),
}
return replica, replica_template
def test_get_artifacts_path_container_with_non_managed_mount_store(self):
store = V1ConnectionType(
name="test_gcs",
kind=V1ConnectionKind.VOLUME_CLAIM,
schema=V1ClaimConnection(mount_path="/claim/path",
volume_claim="claim"),
)
container = get_artifacts_path_container(
polyaxon_init=V1PolyaxonInitContainer(
image="init", image_pull_policy="IfNotPresent"),
artifacts_store=store,
run_path="run_uid",
auto_resume=True,
)
init_args = init_artifact_context_args("run_uid")
init_args.append(
get_volume_args(
store=store,
mount_path=CONTEXT_MOUNT_ARTIFACTS,
artifacts=V1ArtifactsType(dirs=["run_uid"]),
))
assert container == get_base_store_container(
container=k8s_schemas.V1Container(name="init"),
container_name=generate_container_name(
INIT_ARTIFACTS_CONTAINER_PREFIX, "default", False),
polyaxon_init=V1PolyaxonInitContainer(
image="init", image_pull_policy="IfNotPresent"),
store=store,
env=[],
env_from=[],
volume_mounts=[get_artifacts_context_mount()],
args=[" ".join(init_args)],
)
container = get_artifacts_path_container(
polyaxon_init=V1PolyaxonInitContainer(
image="init", image_pull_policy="IfNotPresent"),
artifacts_store=store,
run_path="run_uid",
auto_resume=False,
)
init_args = init_artifact_context_args("run_uid")
assert container == get_base_store_container(
container=k8s_schemas.V1Container(name="init"),
container_name=generate_container_name(
INIT_ARTIFACTS_CONTAINER_PREFIX, "default", False),
polyaxon_init=V1PolyaxonInitContainer(
image="init", image_pull_policy="IfNotPresent"),
store=store,
env=[],
env_from=[],
volume_mounts=[get_artifacts_context_mount()],
args=[" ".join(init_args)],
)
def get_dockerfile_init_container(
polyaxon_init: V1PolyaxonInitContainer,
dockerfile_args: V1DockerfileType,
contexts: PluginsContextsSpec,
run_path: str,
env: List[k8s_schemas.V1EnvVar] = None,
mount_path: Optional[str] = None,
) -> k8s_schemas.V1Container:
env = to_list(env, check_none=True)
env = env + [get_run_instance_env_var()]
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))
return k8s_schemas.V1Container(
name=generate_container_name(INIT_DOCKERFILE_CONTAINER_PREFIX),
image=polyaxon_init.get_image(),
image_pull_policy=polyaxon_init.image_pull_policy,
command=["polyaxon", "docker", "generate"],
args=[
"--build-context={}".format(dockerfile_args.to_dict(dump=True)),
"--destination={}".format(mount_path),
"--copy-path={}".format(
CONTEXT_MOUNT_ARTIFACTS_FORMAT.format(run_path) + "/outputs"),
],
env=env,
resources=polyaxon_init.get_resources(),
volume_mounts=volume_mounts,
)
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 test_config_list_args(self):
config_dict = {
"name": "foo",
"image": "foo/bar:latest",
"args": ["foo"]
}
config = k8s_schemas.V1Container(**config_dict)
assert get_container_command_args(config) == ([], ["foo"])
config_dict = {
"name": "foo",
"image": "foo/bar:latest",
"args": ["foo", "bar"]
}
config = k8s_schemas.V1Container(**config_dict)
assert get_container_command_args(config) == ([], ["foo", "bar"])
def test_get_main_container_with_none_values(self):
container = get_main_container(
container_id="test",
main_container=k8s_schemas.V1Container(name="main"),
contexts=None,
volume_mounts=None,
log_level=None,
artifacts_store=None,
init=None,
connection_by_names=None,
connections=None,
secrets=None,
config_maps=None,
kv_env_vars=None,
env=None,
ports=None,
run_path=None,
)
assert container.name == "test"
assert container.image is None
assert container.image_pull_policy is None
assert container.command is None
assert container.args is None
assert container.ports == []
assert container.env == []
assert container.env_from == []
assert container.resources is None
assert container.volume_mounts == []
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_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 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 test_get_base_store_container_with_store_without_secret(self):
bucket_store_without_secret = V1ConnectionType(
name="test_gcs",
kind=V1ConnectionKind.GCS,
schema=V1BucketConnection(bucket="gs//:foo"),
)
container = get_base_store_container(
container=k8s_schemas.V1Container(name="test"),
container_name="init",
polyaxon_init=V1PolyaxonInitContainer(image_tag=""),
store=bucket_store_without_secret,
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=bucket_store_without_secret, secret=None)
assert container.env_from == []
assert container.resources is not None
assert container.volume_mounts == []
def test_get_artifacts_path_container_with_managed_mount_store(self):
store = V1ConnectionType(
name="test_gcs",
kind=V1ConnectionKind.VOLUME_CLAIM,
schema=V1ClaimConnection(mount_path="/claim/path", volume_claim="claim"),
)
container = get_artifacts_path_container(
polyaxon_init=V1PolyaxonInitContainer(
image="init", image_pull_policy="IfNotPresent"
),
artifacts_store=store,
run_path="run_uid",
clean=True,
)
init_args = init_artifact_context_args("run_uid")
init_args.append(
get_artifacts_store_args(
artifacts_path=get_path(store.store_path, "run_uid"), clean=True
)
)
assert container == get_base_store_container(
container=k8s_schemas.V1Container(name="default"),
container_name=INIT_ARTIFACTS_CONTAINER.format("default"),
polyaxon_init=V1PolyaxonInitContainer(
image="init", image_pull_policy="IfNotPresent"
),
store=store,
env=[],
env_from=[],
volume_mounts=[get_artifacts_context_mount()],
args=[" ".join(init_args)],
is_artifact_store=True,
)
def test_get_artifacts_path_container_with_bucket_store(self):
store = V1ConnectionType(
name="test_gcs",
kind=V1ConnectionKind.GCS,
schema=V1BucketConnection(bucket="gs//:foo"),
)
container = get_artifacts_path_container(
polyaxon_init=V1PolyaxonInitContainer(
image="init", image_pull_policy="IfNotPresent"),
artifacts_store=store,
run_path="run_uid",
auto_resume=True,
)
init_args = init_artifact_context_args("run_uid")
init_args.append(
get_volume_args(
store=store,
mount_path=CONTEXT_MOUNT_ARTIFACTS,
artifacts=V1ArtifactsType(dirs=["run_uid"]),
))
assert container == get_base_store_container(
container=k8s_schemas.V1Container(name="default"),
container_name=INIT_ARTIFACTS_CONTAINER.format("default"),
polyaxon_init=V1PolyaxonInitContainer(
image="init", image_pull_policy="IfNotPresent"),
store=store,
env=[],
env_from=[],
volume_mounts=[get_artifacts_context_mount()],
args=[" ".join(init_args)],
)
def get_store_container(
polyaxon_init: V1PolyaxonInitContainer,
connection: V1ConnectionType,
artifacts: V1ArtifactsType,
container: Optional[k8s_schemas.V1Container] = None,
env: List[k8s_schemas.V1EnvVar] = None,
mount_path: str = None,
) -> Optional[k8s_schemas.V1Container]:
container_name = INIT_ARTIFACTS_CONTAINER.format(connection.name)
if not container:
container = k8s_schemas.V1Container(name=container_name)
volume_name = (get_volume_name(mount_path)
if mount_path else constants.CONTEXT_VOLUME_ARTIFACTS)
mount_path = mount_path or CONTEXT_MOUNT_ARTIFACTS_FORMAT.format(
connection.name)
volume_mounts = [
get_connections_context_mount(name=volume_name, mount_path=mount_path)
]
return get_base_store_container(
container=container,
container_name=container_name,
polyaxon_init=polyaxon_init,
store=connection,
env=env,
env_from=[],
volume_mounts=volume_mounts,
args=[
get_volume_args(store=connection,
mount_path=mount_path,
artifacts=artifacts)
],
)
def test_config_str_command(self):
config_dict = {
"name": "foo",
"image": "foo/bar:latest",
"command": "foo"
}
config = k8s_schemas.V1Container(**config_dict)
assert get_container_command_args(config) == (["foo"], [])
def test_exec_config_with_str_command(self):
config_dict = {
"name": "foo",
"image": "test/test",
"command": "python t2t-trainer",
}
config = k8s_schemas.V1Container(**config_dict)
assert get_container_command_args(config) == (["python t2t-trainer"],
[])
config_dict = {
"name": "foo",
"image": "test/test",
"command": ["python t2t-trainer"],
}
config = k8s_schemas.V1Container(**config_dict)
assert get_container_command_args(config) == (["python t2t-trainer"],
[])
def test_get_main_container_simple_params(self):
initial_mounts = [
k8s_schemas.V1VolumeMount(name="test",
mount_path="/mount_test",
read_only=True)
]
resources = k8s_schemas.V1ResourceRequirements(
requests={
"cpu": "1",
"memory": "256Mi"
},
limits={
"cpu": "1",
"memory": "256Mi"
},
)
container = get_main_container(
container_id="new-name",
main_container=k8s_schemas.V1Container(
name="main",
image="job_docker_image",
image_pull_policy="IfNotPresent",
command=["cmd", "-p", "-c"],
args=["arg1", "arg2"],
resources=resources,
),
contexts=None,
volume_mounts=initial_mounts,
log_level="info",
artifacts_store=None,
init=None,
connection_by_names=None,
connections=None,
secrets=None,
config_maps=None,
kv_env_vars=None,
env=None,
ports=23,
run_path=None,
)
assert container.name == "new-name"
assert container.image == "job_docker_image"
assert container.image_pull_policy == "IfNotPresent"
assert container.command == ["cmd", "-p", "-c"]
assert container.args == ["arg1", "arg2"]
assert container.ports == [
k8s_schemas.V1ContainerPort(container_port=23)
]
assert container.env == [
get_env_var(name=POLYAXON_KEYS_LOG_LEVEL, value="info")
]
assert container.env_from == []
assert container.resources == resources
assert container.volume_mounts == initial_mounts
def test_main_container(self):
store = V1ConnectionType(
name="test_gcs",
kind=V1ConnectionKind.S3,
schema=V1BucketConnection(bucket="s3://foo"),
secret=None,
)
contexts = PluginsContextsSpec.from_config(
V1Plugins.from_dict({}), default_auth=True
)
main_container = k8s_schemas.V1Container(
name="main",
image="foo/test",
image_pull_policy="IfNotPresent",
command=["foo", "bar"],
args=["arg1", "arg2"],
)
container = self.converter.get_main_container(
main_container=main_container,
contexts=contexts,
artifacts_store=store,
init_connections=[],
connections=[],
connection_by_names={},
log_level="info",
secrets=[],
config_maps=[],
kv_env_vars=[],
ports=None,
)
expected_container = get_main_container(
container_id="dummy",
main_container=main_container,
contexts=contexts,
volume_mounts=get_mounts(
use_auth_context=True,
use_artifacts_context=False,
use_docker_context=False,
use_shm_context=False,
),
log_level="info",
artifacts_store=store,
connections=[],
init=[],
connection_by_names={},
secrets=[],
config_maps=[],
kv_env_vars=[],
env=self.converter.get_main_env_vars(),
ports=None,
run_path="/test",
)
assert container == expected_container
def test_get_service_custom_resource_missing_keys(self):
main_container = k8s_schemas.V1Container(name="main")
metadata, pod_spec = get_pod_spec(
namespace="default",
main_container=main_container,
sidecar_containers=None,
init_containers=None,
resource_name="foo",
volumes=[],
environment=None,
labels=None,
annotations=None,
)
notifications = [
V1Notification(connections=["test"], trigger=V1Statuses.DONE)
]
custom_object = {
"template": get_pod_template_spec(metadata=metadata,
pod_spec=pod_spec),
"ports": [12, 121, 12],
}
expected_crd = get_custom_object(
namespace="default",
resource_name="foo",
kind="Operation",
api_version="core.polyaxon.com/v1",
labels=None,
annotations=None,
custom_object={
"serviceSpec": custom_object,
"collectLogs": False,
"syncStatuses": False,
"notifications": [n.to_operator() for n in notifications],
},
)
crd = get_service_custom_resource(
namespace="default",
resource_name="foo",
main_container=main_container,
sidecar_containers=None,
init_containers=None,
volumes=[],
termination=None,
collect_logs=None,
sync_statuses=None,
notifications=notifications,
environment=None,
labels=None,
annotations=None,
ports=[12, 121, 12],
)
assert crd == expected_crd
def test_get_base_store_container_with_none_values(self):
with self.assertRaises(PolypodException):
get_base_store_container(
container=k8s_schemas.V1Container(name="init"),
container_name=None,
polyaxon_init=V1PolyaxonInitContainer(),
store=None,
env=None,
env_from=None,
volume_mounts=None,
args=None,
)
def get_auth_context_container(
polyaxon_init: V1PolyaxonInitContainer, env: List[k8s_schemas.V1EnvVar] = None
) -> k8s_schemas.V1Container:
env = to_list(env, check_none=True)
return k8s_schemas.V1Container(
name=INIT_AUTH_CONTAINER,
image=polyaxon_init.get_image(),
image_pull_policy=polyaxon_init.image_pull_policy,
command=["polyaxon", "initializer", "auth"],
env=env,
resources=polyaxon_init.get_resources(),
volume_mounts=[get_auth_context_mount(read_only=False)],
)
def test_get_main_container_host_paths(self):
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,
)
volume_mounts = get_mounts(
use_auth_context=contexts.auth,
use_artifacts_context=False,
use_docker_context=contexts.docker,
use_shm_context=contexts.shm,
)
artifacts_store = V1ConnectionType(
name="plx-outputs",
kind=V1ConnectionKind.HOST_PATH,
schema=V1HostPathConnection(mount_path="/tmp/plx/outputs",
host_path="/tmp/plx/outputs"),
)
container = get_main_container(
container_id="test",
main_container=k8s_schemas.V1Container(name="main"),
contexts=PluginsContextsSpec.from_config(
V1Plugins(collect_artifacts=True, collect_logs=True)),
volume_mounts=volume_mounts,
log_level=None,
artifacts_store=artifacts_store,
init=[],
connections=[],
connection_by_names={artifacts_store.name: artifacts_store},
secrets=[],
config_maps=[],
kv_env_vars=None,
env=None,
ports=None,
run_path="run_path",
)
assert container.volume_mounts == [
get_auth_context_mount(read_only=True),
get_artifacts_context_mount(read_only=False),
]
def test_get_main_container(self):
container = get_main_container(
container_id="test",
main_container=k8s_schemas.V1Container(name="main"),
contexts=None,
volume_mounts=None,
log_level=None,
artifacts_store=None,
init=[
V1Init(connection=self.claim_store.name),
V1Init(connection=self.s3_store.name),
],
connections=[self.host_path_store.name, self.gcs_store.name],
connection_by_names={
self.claim_store.name: self.claim_store,
self.s3_store.name: self.s3_store,
self.host_path_store.name: self.host_path_store,
self.gcs_store.name: self.gcs_store,
},
secrets=[self.mount_resource1, self.request_non_mount_resource1],
config_maps=[
self.non_mount_resource1, self.request_mount_resource2
],
kv_env_vars=None,
env=None,
ports=None,
run_path="run_path",
)
assert container.name == "test"
assert container.image is None
assert container.image_pull_policy is None
assert container.command is None
assert container.args is None
assert container.ports == []
# 2 env vars from the secret mount
# + 2 for the connection (context_path + spec)
# + 1 for the connection spec (non mount)
assert len(container.env) == 5
assert container.env_from == []
assert container.resources is None
assert len(container.volume_mounts) == 4
def get_file_init_container(
polyaxon_init: V1PolyaxonInitContainer,
file_args: V1FileType,
contexts: PluginsContextsSpec,
run_path: str,
run_instance: str,
env: List[k8s_schemas.V1EnvVar] = None,
mount_path: Optional[str] = None,
) -> k8s_schemas.V1Container:
env = to_list(env, check_none=True)
env = env + [get_run_instance_env_var(run_instance)]
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))
file_args.filename = file_args.filename or "file"
return k8s_schemas.V1Container(
name=generate_container_name(INIT_FILE_CONTAINER_PREFIX),
image=polyaxon_init.get_image(),
image_pull_policy=polyaxon_init.image_pull_policy,
command=["polyaxon", "initializer", "file"],
args=[
"--file-context={}".format(file_args.to_dict(dump=True)),
"--filepath={}".format(mount_path),
"--copy-path={}".format(
CONTEXT_MOUNT_RUN_OUTPUTS_FORMAT.format(run_path)),
"--track",
],
env=env,
resources=polyaxon_init.get_resources(),
volume_mounts=volume_mounts,
)
def test_get_main_container_with_bucket_artifacts_store(self):
container = get_main_container(
container_id="main",
main_container=k8s_schemas.V1Container(name="main"),
contexts=PluginsContextsSpec.from_config(
V1Plugins(collect_artifacts=True,
collect_logs=True,
collect_resources=True)),
volume_mounts=None,
log_level=None,
artifacts_store=self.s3_store,
init=None,
connections=None,
connection_by_names={self.s3_store.name: self.s3_store},
secrets=None,
config_maps=None,
kv_env_vars=None,
env=None,
ports=None,
run_path="run_path",
)
assert container.name == "main"
assert container.image is None
assert container.image_pull_policy is None
assert container.command is None
assert container.args is None
assert container.ports == []
assert len(container.env) == 2
assert container.env_from == []
assert container.resources is None
assert len(container.volume_mounts) == 1 # mount context
container = get_main_container(
container_id="main1",
main_container=k8s_schemas.V1Container(name="main"),
contexts=PluginsContextsSpec.from_config(
V1Plugins(
collect_artifacts=True,
collect_logs=True,
collect_resources=True,
sync_statuses=True,
)),
volume_mounts=None,
log_level=None,
artifacts_store=self.s3_store,
init=None,
connections=None,
connection_by_names={self.s3_store.name: self.s3_store},
secrets=[self.mount_resource1],
config_maps=None,
kv_env_vars=None,
env=None,
ports=None,
run_path="run_path",
)
assert container.name == "main1"
assert container.image is None
assert container.image_pull_policy is None
assert container.command is None
assert container.args is None
assert container.ports == []
assert len(container.env) == 2
assert container.env_from == []
assert container.resources is None
# The mount resource1 is not requested
assert len(container.volume_mounts) == 1 # one mount resource
container = get_main_container(
container_id="main1",
main_container=k8s_schemas.V1Container(name="main"),
contexts=PluginsContextsSpec.from_config(
V1Plugins(collect_artifacts=True,
collect_logs=True,
collect_resources=True)),
volume_mounts=None,
log_level=None,
artifacts_store=self.s3_store,
init=None,
connections=None,
connection_by_names={self.s3_store.name: self.s3_store},
secrets=[self.request_mount_resource2],
config_maps=None,
kv_env_vars=None,
env=None,
ports=None,
run_path="run_path",
)
assert container.name == "main1"
assert container.image is None
assert container.image_pull_policy is None
assert container.command is None
assert container.args is None
assert container.ports == []
assert len(container.env) == 2
assert container.env_from == []
assert container.resources is None
# The mount resource2 is requested
assert len(container.volume_mounts) == 2 # one mount resource
container = get_main_container(
container_id="tensorflow",
main_container=k8s_schemas.V1Container(name="main"),
contexts=PluginsContextsSpec.from_config(
V1Plugins(collect_artifacts=True,
collect_logs=True,
collect_resources=False)),
volume_mounts=None,
log_level=None,
artifacts_store=self.s3_store,
init=None,
connections=None,
connection_by_names={self.s3_store.name: self.s3_store},
secrets=[self.non_mount_resource1],
config_maps=None,
kv_env_vars=None,
env=None,
ports=None,
run_path="run_path",
)
assert container.name == "tensorflow"
assert container.image is None
assert container.image_pull_policy is None
assert container.command is None
assert container.args is None
assert container.ports == []
assert len(container.env) == 1
assert container.env_from == []
assert container.resources is None
assert len(container.volume_mounts) == 1 # outputs context
container = get_main_container(
container_id="pytorch",
main_container=k8s_schemas.V1Container(name="main"),
contexts=PluginsContextsSpec.from_config(
V1Plugins(collect_artifacts=True,
collect_logs=True,
collect_resources=True)),
volume_mounts=None,
log_level=None,
artifacts_store=self.s3_store,
init=None,
connections=None,
connection_by_names={self.s3_store.name: self.s3_store},
secrets=[self.request_non_mount_resource1],
config_maps=None,
kv_env_vars=None,
env=None,
ports=None,
run_path="run_path",
)
assert container.name == "pytorch"
assert container.image is None
assert container.image_pull_policy is None
assert container.command is None
assert container.args is None
assert container.ports == []
assert len(
container.env) == 2 + 2 # 2 + 2 env vars from the secret mount
assert container.env_from == []
assert container.resources is None
assert len(container.volume_mounts) == 1
def test_get_main_container_with_mounted_artifacts_store(self):
container = get_main_container(
container_id="test",
main_container=k8s_schemas.V1Container(name="main"),
contexts=None,
volume_mounts=None,
log_level=None,
artifacts_store=None,
init=[V1Init(connection=self.claim_store.name)],
connections=None,
connection_by_names={self.claim_store.name: self.claim_store},
secrets=None,
config_maps=None,
kv_env_vars=None,
env=None,
ports=None,
run_path="run_path",
)
assert container.name == "test"
assert container.image is None
assert container.image_pull_policy is None
assert container.command is None
assert container.args is None
assert container.ports == []
assert container.env_from == []
assert container.resources is None
assert len(container.volume_mounts) == 1
container = get_main_container(
container_id="",
main_container=k8s_schemas.V1Container(name="main"),
contexts=None,
volume_mounts=None,
log_level=None,
artifacts_store=None,
init=[V1Init(connection=self.claim_store.name)],
connections=[self.claim_store.name],
connection_by_names={self.claim_store.name: self.claim_store},
secrets=None,
config_maps=None,
kv_env_vars=None,
env=None,
ports=None,
run_path="run_path",
)
assert container.name == "main"
assert container.image is None
assert container.image_pull_policy is None
assert container.command is None
assert container.args is None
assert container.ports == []
assert container.env_from == []
assert container.resources is None
assert len(container.volume_mounts) == 2
container = get_main_container(
container_id="main-job",
main_container=k8s_schemas.V1Container(name="main"),
contexts=PluginsContextsSpec.from_config(
V1Plugins(collect_artifacts=True,
collect_logs=True,
collect_resources=True)),
volume_mounts=None,
log_level=None,
artifacts_store=self.claim_store,
init=None,
connections=[],
connection_by_names={self.claim_store.name: self.claim_store},
secrets=None,
config_maps=None,
kv_env_vars=None,
env=None,
ports=None,
run_path="run_path",
)
assert container.name == "main-job"
assert container.image is None
assert container.image_pull_policy is None
assert container.command is None
assert container.args is None
assert container.ports == []
assert len(container.env) == 2
assert container.env_from == []
assert container.resources is None
assert len(container.volume_mounts) == 1
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(),
)
def test_get_pod_spec(self):
init_container = k8s_schemas.V1Container(name="init")
main_container = k8s_schemas.V1Container(name="main")
sidecar_container = k8s_schemas.V1Container(name="sidecar")
volumes = [k8s_schemas.V1Volume(name="vol")]
labels = {"key": "labels"}
annotations = {"key": "annotations"}
node_selector = {"key": "selector"}
affinity = [{"key": "affinity"}]
tolerations = {"key": "tolerations"}
security_context = {"uid": 222, "gid": 222}
restart_policy = "never"
with self.assertRaises(PolypodException):
get_pod_spec(
namespace="default",
main_container=None,
sidecar_containers=None,
init_containers=None,
resource_name="foo",
volumes=None,
environment=V1Environment(),
labels={},
annotations={},
)
environment = V1Environment(
service_account_name="sa",
labels=labels,
annotations=annotations,
node_selector=node_selector,
affinity=affinity,
tolerations=tolerations,
security_context=security_context,
image_pull_secrets=[],
restart_policy=restart_policy,
)
metadata, pod_spec = get_pod_spec(
namespace="default",
resource_name="foo",
main_container=main_container,
sidecar_containers=None,
init_containers=None,
volumes=None,
environment=environment,
labels=environment.labels,
annotations=environment.annotations,
)
assert metadata.name == "foo"
assert metadata.labels == labels
assert metadata.namespace == "default"
assert metadata.annotations == annotations
assert isinstance(pod_spec, k8s_schemas.V1PodSpec)
assert pod_spec.security_context == security_context
assert pod_spec.restart_policy == "never"
assert pod_spec.service_account_name == "sa"
assert pod_spec.init_containers == []
assert pod_spec.containers == [main_container]
assert pod_spec.volumes is None
assert pod_spec.node_selector == node_selector
assert pod_spec.tolerations == tolerations
assert pod_spec.affinity == affinity
environment = V1Environment(
service_account_name="sa",
labels=labels,
annotations=annotations,
node_selector=node_selector,
affinity=affinity,
tolerations=tolerations,
security_context=security_context,
image_pull_secrets=[],
restart_policy=restart_policy,
)
metadata, pod_spec = get_pod_spec(
namespace="default",
main_container=main_container,
sidecar_containers=[sidecar_container],
init_containers=[init_container],
resource_name="foo",
volumes=volumes,
environment=environment,
labels={},
annotations={},
)
assert pod_spec.init_containers == [init_container]
assert pod_spec.containers == [main_container, sidecar_container]
assert pod_spec.volumes == volumes
assert metadata.annotations == {}
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,
)
