def test_get_file_init_container(self):
file_args = V1FileType(content="test")
container = get_file_init_container(
polyaxon_init=V1PolyaxonInitContainer(image="foo", image_tag=""),
contexts=PluginsContextsSpec.from_config(V1Plugins(auth=True)),
file_args=V1FileType(content="test"),
run_path="test",
run_instance="foo.bar.runs.uuid",
)
assert INIT_FILE_CONTAINER_PREFIX in container.name
assert container.image == "foo"
assert container.image_pull_policy is None
assert container.command == ["polyaxon", "initializer", "file"]
assert container.resources == get_init_resources()
assert container.volume_mounts == [
get_connections_context_mount(
name=constants.CONTEXT_VOLUME_ARTIFACTS,
mount_path=CONTEXT_MOUNT_ARTIFACTS,
),
get_auth_context_mount(read_only=True),
]
assert file_args.to_dict(dump=True) == '{"content":"test"}'
assert container.args == [
"--file-context={}".format('{"content":"test","filename":"file"}'),
"--filepath={}".format(CONTEXT_MOUNT_ARTIFACTS),
"--copy-path={}".format(
CONTEXT_MOUNT_RUN_OUTPUTS_FORMAT.format("test")),
"--track",
]
file_args = V1FileType(filename="test", content="test")
container = get_file_init_container(
polyaxon_init=V1PolyaxonInitContainer(
image="init/init",
image_tag="",
image_pull_policy="IfNotPresent"),
contexts=PluginsContextsSpec.from_config(V1Plugins(auth=True)),
file_args=file_args,
run_path="test",
run_instance="foo.bar.runs.uuid",
)
assert INIT_FILE_CONTAINER_PREFIX in container.name
assert container.image == "init/init"
assert container.image_pull_policy == "IfNotPresent"
assert container.command == ["polyaxon", "initializer", "file"]
assert container.args == [
"--file-context={}".format(file_args.to_dict(dump=True)),
"--filepath={}".format(CONTEXT_MOUNT_ARTIFACTS),
"--copy-path={}".format(
CONTEXT_MOUNT_RUN_OUTPUTS_FORMAT.format("test")),
"--track",
]
assert container.resources == get_init_resources()
assert container.volume_mounts == [
get_connections_context_mount(
name=constants.CONTEXT_VOLUME_ARTIFACTS,
mount_path=CONTEXT_MOUNT_ARTIFACTS,
),
get_auth_context_mount(read_only=True),
]
def test_get_env_vars_with_artifacts_store(self):
assert (get_env_vars(
contexts=None,
log_level=None,
kv_env_vars=None,
connections=None,
secrets=None,
config_maps=None,
) == [])
assert get_env_vars(
contexts=PluginsContextsSpec.from_config(
V1Plugins(collect_logs=False,
collect_artifacts=True,
collect_resources=True)),
log_level=None,
kv_env_vars=None,
connections=None,
secrets=None,
config_maps=None,
) == [
get_env_var(name=POLYAXON_KEYS_COLLECT_ARTIFACTS, value=True),
get_env_var(name=POLYAXON_KEYS_COLLECT_RESOURCES, value=True),
]
assert (get_env_vars(
contexts=PluginsContextsSpec.from_config(
V1Plugins(
collect_logs=False,
collect_artifacts=False,
collect_resources=False,
)),
log_level=None,
kv_env_vars=None,
connections=None,
secrets=None,
config_maps=None,
) == [])
assert (get_env_vars(
contexts=None,
log_level=None,
kv_env_vars=None,
connections=None,
secrets=None,
config_maps=None,
) == [])
assert get_env_vars(
contexts=PluginsContextsSpec.from_config(
V1Plugins(collect_logs=False,
collect_artifacts=True,
collect_resources=False)),
log_level=None,
kv_env_vars=None,
connections=None,
secrets=None,
config_maps=None,
) == [get_env_var(name=POLYAXON_KEYS_COLLECT_ARTIFACTS, value=True)]
def test_default_volumes(self):
assert (get_pod_volumes(
contexts=None,
artifacts_store=None,
init_connections=None,
connection_by_names=None,
secrets=None,
config_maps=None,
volumes=None,
) == [])
assert (get_pod_volumes(
contexts=None,
artifacts_store=None,
init_connections=[],
connection_by_names={},
secrets=[],
config_maps=[],
volumes=[],
) == [])
assert (get_pod_volumes(
contexts=PluginsContextsSpec.from_config(
V1Plugins(
docker=False,
shm=False,
auth=False,
collect_artifacts=False,
collect_logs=False,
)),
artifacts_store=None,
init_connections=[],
connection_by_names={},
secrets=[],
config_maps=[],
volumes=[],
) == [])
assert get_pod_volumes(
contexts=PluginsContextsSpec.from_config(
V1Plugins(
docker=True,
shm=True,
auth=True,
collect_artifacts=False,
collect_logs=False,
)),
artifacts_store=None,
init_connections=[],
connection_by_names={},
secrets=[],
config_maps=[],
volumes=[],
) == [
get_shm_context_volume(),
get_configs_context_volume(),
get_docker_context_volume(),
]
def test_get_from_empty_env(self):
spec = PluginsContextsSpec.from_config(V1Plugins(), default_auth=True)
assert spec.auth is True
assert spec.docker is False
assert spec.shm is True
assert spec.collect_artifacts is True
assert spec.collect_logs is True
assert spec.sync_statuses is True
spec = PluginsContextsSpec.from_config(V1Plugins())
assert spec.auth is False
assert spec.docker is False
assert spec.shm is True
assert spec.collect_artifacts is True
assert spec.collect_logs is True
assert spec.sync_statuses is True
def test_get_sidecar_container_with_non_managed_mount_outputs_logs_store(
self):
env_vars = [
get_env_var(name="key1", value="value1"),
get_env_var(name="key2", value="value2"),
]
mount_non_managed_store = V1ConnectionType(
name="test_claim",
kind=V1ConnectionKind.VOLUME_CLAIM,
schema=V1ClaimConnection(volume_claim="test",
mount_path="/tmp",
read_only=True),
)
sidecar = get_sidecar_container(
container_id=MAIN_JOB_CONTAINER,
env=env_vars,
polyaxon_sidecar=V1PolyaxonSidecarContainer(
image="sidecar/sidecar",
image_pull_policy="IfNotPresent",
sleep_interval=213,
sync_interval=-1,
),
artifacts_store=mount_non_managed_store,
contexts=PluginsContextsSpec.from_config(
V1Plugins(collect_logs=False,
collect_artifacts=False,
auth=True)),
run_path=None,
)
assert sidecar is None
def test_get_sidecar_container_with_non_managed_bucket_artifacts_logs_store(
self):
env_vars = [
get_env_var(name="key1", value="value1"),
get_env_var(name="key2", value="value2"),
]
bucket_non_managed_store = V1ConnectionType(
name="test_s3",
kind=V1ConnectionKind.S3,
schema=V1BucketConnection(bucket="s3//:foo"),
)
sidecar = get_sidecar_container(
container_id=MAIN_JOB_CONTAINER,
env=env_vars,
polyaxon_sidecar=V1PolyaxonSidecarContainer(
image="sidecar/sidecar",
image_pull_policy="IfNotPresent",
sleep_interval=213,
sync_interval=-1,
),
artifacts_store=bucket_non_managed_store,
contexts=PluginsContextsSpec.from_config(
V1Plugins(collect_logs=False,
collect_artifacts=False,
auth=True)),
run_path=None,
)
assert sidecar is None
def get_resource(
self,
compiled_operation: V1CompiledOperation,
artifacts_store: V1ConnectionType,
connection_by_names: Dict[str, V1ConnectionType],
secrets: Optional[Iterable[V1K8sResourceType]],
config_maps: Optional[Iterable[V1K8sResourceType]],
default_sa: str = None,
default_auth: bool = False,
) -> Dict:
job = compiled_operation.run # type: V1TFJob
def _get_replica(
replica: Optional[V1KFReplica]) -> Optional[ReplicaSpec]:
if not replica:
return None
return self.get_replica_resource(
plugins=plugins,
contexts=contexts,
environment=replica.environment,
volumes=replica.volumes or [],
init=replica.init or [],
sidecars=replica.sidecars or [],
container=replica.container,
artifacts_store=artifacts_store,
connections=replica.connections or [],
connection_by_names=connection_by_names,
secrets=secrets,
config_maps=config_maps,
default_sa=default_sa,
num_replicas=replica.replicas,
)
plugins = compiled_operation.plugins or V1Plugins()
contexts = PluginsContextsSpec.from_config(plugins,
default_auth=default_auth)
chief = _get_replica(job.chief)
worker = _get_replica(job.worker)
ps = _get_replica(job.ps)
evaluator = _get_replica(job.evaluator)
labels = self.get_labels(version=pkg.VERSION, labels={})
return get_tf_job_custom_resource(
namespace=self.namespace,
resource_name=self.resource_name,
chief=chief,
worker=worker,
ps=ps,
evaluator=evaluator,
termination=compiled_operation.termination,
collect_logs=contexts.collect_logs,
clean_pod_policy=job.clean_pod_policy,
sync_statuses=contexts.sync_statuses,
notifications=plugins.notifications,
labels=labels,
annotations=self.annotations,
)
def test_get_env_vars_with_all(self):
connection = V1ConnectionType(
name="test_s3",
kind=V1ConnectionKind.S3,
schema=V1BucketConnection(bucket="s3//:foo"),
secret=self.resource6.schema,
)
env_vars = get_env_vars(
contexts=PluginsContextsSpec.from_config(
V1Plugins(collect_logs=False,
collect_artifacts=True,
collect_resources=True)),
log_level="info",
kv_env_vars=[["key1", "val1"], ["key2", "val2"]],
artifacts_store_name="test",
connections=[connection],
secrets=[
self.resource1,
self.resource2,
self.resource3,
self.resource4,
self.resource6,
],
config_maps=[
self.resource1,
self.resource2,
self.resource3,
self.resource4,
],
)
expected = [
get_env_var(name=POLYAXON_KEYS_LOG_LEVEL, value="info"),
get_env_var(name=POLYAXON_KEYS_COLLECT_ARTIFACTS, value=True),
get_env_var(name=POLYAXON_KEYS_COLLECT_RESOURCES, value=True),
get_env_var(name=POLYAXON_KEYS_ARTIFACTS_STORE_NAME, value="test"),
]
expected += get_connection_env_var(connection=connection,
secret=self.resource6)
expected += get_kv_env_vars([["key1", "val1"], ["key2", "val2"]])
expected += get_env_vars_from_k8s_resources(
secrets=[
self.resource1,
self.resource2,
self.resource3,
self.resource4,
self.resource6,
],
config_maps=[
self.resource1,
self.resource2,
self.resource3,
self.resource4,
],
)
assert env_vars == expected
def get_tuner(
name: str,
container: V1Container,
matrix: V1Matrix,
search: V1ParamSearch,
iteration: int,
bracket_iteration: int = None,
) -> V1Operation:
params = {
"matrix": V1Param(value=matrix.to_dict()),
"search": V1Param(value=search.to_dict()),
"iteration": V1Param(value=iteration),
}
inputs = [
V1IO(name="matrix", iotype=types.DICT, is_list=False,
is_optional=True),
V1IO(name="search", iotype=types.DICT, is_list=False,
is_optional=True),
V1IO(name="iteration",
iotype=types.INT,
is_list=False,
is_optional=True),
]
if bracket_iteration is not None:
params["bracket_iteration"] = V1Param(value=bracket_iteration)
inputs.append(
V1IO(
name="bracket_iteration",
iotype=types.INT,
is_list=False,
is_optional=True,
))
return V1Operation(
params=params,
component=V1Component(
name=name,
plugins=V1Plugins(
auth=True,
collect_logs=True,
collect_artifacts=True,
collect_resources=False,
sync_statuses=False,
),
inputs=inputs,
outputs=[
V1IO(
name="suggestions",
iotype=types.DICT,
is_list=True,
is_optional=False,
),
],
run=V1Tuner(container=container, ),
),
)
def test_artifacts_store(self):
self.assert_contexts_store(contexts=None, results=[])
self.assert_contexts_store(
contexts=PluginsContextsSpec.from_config(
PluginsContextsSpec.from_config(
V1Plugins(collect_logs=False, collect_artifacts=True))),
results=[get_artifacts_context_mount(read_only=False)],
)
self.assert_contexts_store(
contexts=PluginsContextsSpec.from_config(
PluginsContextsSpec.from_config(
V1Plugins(collect_logs=True, collect_artifacts=False))),
results=[],
)
self.assert_contexts_store(
contexts=PluginsContextsSpec.from_config(
PluginsContextsSpec.from_config(
V1Plugins(collect_logs=True, collect_artifacts=True))),
results=[get_artifacts_context_mount(read_only=False)],
)
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_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 get_resource(
self,
compiled_operation: V1CompiledOperation,
artifacts_store: V1ConnectionType,
connection_by_names: Dict[str, V1ConnectionType],
secrets: Optional[Iterable[V1K8sResourceType]],
config_maps: Optional[Iterable[V1K8sResourceType]],
is_resume: bool = False,
) -> Dict:
job = compiled_operation.run # type: V1MPIJob
def _get_replica(
replica: Optional[V1KFReplica]) -> Optional[ReplicaSpec]:
if not replica:
return None
return self.get_replica_resource(
plugins=plugins,
contexts=contexts,
environment=replica.environment,
volumes=replica.volumes or [],
init=replica.init or [],
sidecars=replica.sidecars or [],
container=replica.container,
artifacts_store=artifacts_store,
connections=replica.connections or [],
connection_by_names=connection_by_names,
secrets=secrets,
config_maps=config_maps,
is_resume=is_resume,
num_replicas=replica.replicas,
)
plugins = compiled_operation.plugins or V1Plugins()
contexts = PluginsContextsSpec.from_config(plugins)
launcher = _get_replica(job.launcher)
worker = _get_replica(job.worker)
labels = self.get_labels(version=pkg.VERSION, labels={})
return get_mpi_job_custom_resource(
namespace=self.namespace,
resource_name=self.resource_name,
launcher=launcher,
worker=worker,
clean_pod_policy=job.clean_pod_policy,
slots_per_worker=job.slots_per_worker,
termination=compiled_operation.termination,
collect_logs=contexts.collect_logs,
sync_statuses=contexts.sync_statuses,
notifications=plugins.notifications,
labels=labels,
)
def get_tuner(
name: str,
container: V1Container,
matrix: V1Matrix,
configs: List[Dict],
metrics: List[float],
iteration: int,
) -> V1Operation:
return V1Operation(
params={
"configs": V1Param(value=configs),
"metrics": V1Param(value=metrics),
"matrix": V1Param(value=matrix),
"iteration": V1Param(value=iteration),
},
termination=V1Termination(max_retries=3),
component=V1Component(
name=name,
plugins=V1Plugins(
auth=True,
collect_logs=False,
collect_artifacts=False,
collect_resources=False,
sync_statuses=True,
),
inputs=[
V1IO(name="configs",
iotype=types.DICT,
is_list=True,
is_optional=False),
V1IO(name="metrics",
iotype=types.FLOAT,
is_list=True,
is_optional=False),
V1IO(name="iteration",
iotype=types.INT,
is_list=True,
is_optional=True),
],
outputs=[
V1IO(
name="suggestions",
iotype=types.DICT,
is_list=True,
is_optional=False,
),
],
run=V1Tuner(container=container, ),
),
)
def test_get_from_env(self):
config = V1Plugins(
auth=True,
shm=True,
docker=True,
collect_artifacts=True,
collect_logs=True,
sync_statuses=True,
)
spec = PluginsContextsSpec.from_config(config)
assert spec.auth is True
assert spec.docker is True
assert spec.shm is True
assert spec.collect_artifacts is True
assert spec.collect_logs is True
assert spec.sync_statuses is True
def assert_artifacts_store_raises(self, store, run_path=None):
with self.assertRaises(PolypodException):
get_sidecar_container(
container_id=MAIN_JOB_CONTAINER,
contexts=PluginsContextsSpec.from_config(
V1Plugins(collect_logs=True, collect_artifacts=False)),
env=None,
polyaxon_sidecar=V1PolyaxonSidecarContainer(
image="sidecar/sidecar",
image_pull_policy="IfNotPresent",
sleep_interval=213,
sync_interval=213,
),
artifacts_store=store,
run_path=run_path,
)
def assert_artifacts_store(store, results):
assert (get_pod_volumes(
contexts=PluginsContextsSpec.from_config(
V1Plugins(
docker=False,
shm=False,
auth=False,
collect_artifacts=True,
collect_logs=False,
)),
artifacts_store=store,
init_connections=[],
connection_by_names={},
secrets=[],
config_maps=[],
volumes=[],
) == results)
def test_get_init_containers_with_auth(self):
containers = self.converter.get_init_containers(
polyaxon_init=V1PolyaxonInitContainer(image="foo/foo"),
contexts=PluginsContextsSpec.from_config(
V1Plugins(collect_logs=False, collect_artifacts=False)),
artifacts_store=None,
init_connections=None,
is_resume=False,
connection_by_names={},
init_containers=[],
)
assert containers == [
get_auth_context_container(
polyaxon_init=V1PolyaxonInitContainer(image="foo/foo"),
env=self.converter.get_auth_service_env_vars(),
)
]
def get_resource(
self,
compiled_operation: V1CompiledOperation,
artifacts_store: V1ConnectionType,
connection_by_names: Dict[str, V1ConnectionType],
secrets: Optional[Iterable[V1K8sResourceType]],
config_maps: Optional[Iterable[V1K8sResourceType]],
default_sa: str = None,
default_auth: bool = False,
) -> Dict:
service = compiled_operation.run # type: V1Service
plugins = compiled_operation.plugins or V1Plugins()
contexts = PluginsContextsSpec.from_config(plugins,
default_auth=default_auth)
replica_spec = self.get_replica_resource(
plugins=plugins,
contexts=contexts,
environment=service.environment,
volumes=service.volumes,
init=service.init,
sidecars=service.sidecars,
container=service.container,
artifacts_store=artifacts_store,
connections=service.connections,
connection_by_names=connection_by_names,
secrets=secrets,
config_maps=config_maps,
default_sa=default_sa,
ports=service.ports,
)
return get_service_custom_resource(
namespace=self.namespace,
main_container=replica_spec.main_container,
sidecar_containers=replica_spec.sidecar_containers,
init_containers=replica_spec.init_containers,
resource_name=self.resource_name,
volumes=replica_spec.volumes,
environment=replica_spec.environment,
termination=compiled_operation.termination,
collect_logs=contexts.collect_logs,
sync_statuses=contexts.sync_statuses,
notifications=plugins.notifications,
labels=replica_spec.labels,
annotations=self.annotations,
ports=service.ports,
)
def test_passing_volumes(self):
assert get_pod_volumes(
contexts=PluginsContextsSpec.from_config(
V1Plugins(
docker=False,
shm=False,
auth=False,
collect_artifacts=False,
collect_logs=False,
)),
artifacts_store=None,
init_connections=[],
connection_by_names={},
secrets=[],
config_maps=[],
volumes=[self.vol1, self.vol2, self.vol3],
) == [self.vol1, self.vol2, self.vol3]
def get_notifier_operation(
connection: str,
backend: str,
owner: str,
project: str,
run_uuid: str,
run_name: str,
condition: Union[str, Dict],
) -> V1Operation:
return V1Operation(
params={
"backend": V1Param(value=backend),
"owner": V1Param(value=owner),
"project": V1Param(value=project),
"uuid": V1Param(value=run_uuid),
"name": V1Param(value=run_name),
"condition": V1Param(value=condition),
},
termination=V1Termination(max_retries=3),
component=V1Component(
name="notifier",
plugins=V1Plugins(
auth=False,
collect_logs=False,
collect_artifacts=False,
collect_resources=False,
auto_resume=False,
sync_statuses=False,
external_host=True,
),
inputs=[
V1IO(name="backend", type=types.STR, is_optional=False),
V1IO(name="owner", type=types.STR, is_optional=False),
V1IO(name="project", type=types.STR, is_optional=False),
V1IO(name="uuid", type=types.STR, is_optional=False),
V1IO(name="name", type=types.STR, is_optional=True),
V1IO(name="condition", type=types.DICT, is_optional=True),
V1IO(name="connection", type=types.STR, is_optional=True),
],
run=V1NotifierJob(
connections=[connection],
container=get_default_notification_container(),
),
),
)
def get_notifier_operation(
connection: str,
kind: str,
owner: str,
project: str,
run_uuid: str,
run_name: str,
condition: str,
) -> V1Operation:
return V1Operation(
params={
"kind": V1Param(value=kind),
"owner": V1Param(value=owner),
"project": V1Param(value=project),
"run_uuid": V1Param(value=run_uuid),
"run_name": V1Param(value=run_name),
"condition": V1Param(value=condition),
},
termination=V1Termination(max_retries=3),
component=V1Component(
name="slack-notification",
plugins=V1Plugins(
auth=False,
collect_logs=False,
collect_artifacts=False,
collect_resources=False,
sync_statuses=False,
),
inputs=[
V1IO(name="kind", iotype=types.STR, is_optional=False),
V1IO(name="owner", iotype=types.STR, is_optional=False),
V1IO(name="project", iotype=types.STR, is_optional=False),
V1IO(name="run_uuid", iotype=types.STR, is_optional=False),
V1IO(name="run_name", iotype=types.STR, is_optional=True),
V1IO(name="condition", iotype=types.STR, is_optional=True),
V1IO(name="connection", iotype=types.STR, is_optional=True),
],
run=V1Notifier(
connections=[connection],
container=get_default_notification_container(),
),
),
)
def assert_artifacts_store_raises(self, store, run_path):
with self.assertRaises(PolypodException):
get_main_container(
main_container=None,
contexts=PluginsContextsSpec.from_config(
V1Plugins(collect_artifacts=True, collect_logs=False)),
volume_mounts=None,
log_level=None,
artifacts_store=store,
init=None,
connection_by_names=None,
connections=None,
secrets=None,
config_maps=None,
kv_env_vars=None,
env=None,
ports=None,
run_path=run_path,
)
def test_get_init_containers_with_store_outputs(self):
store = V1ConnectionType(
name="test_gcs",
kind=V1ConnectionKind.S3,
schema=V1BucketConnection(bucket="s3://foo"),
)
# No context
containers = self.converter.get_init_containers(
contexts=None,
artifacts_store=store,
init_connections=[],
init_containers=[],
connection_by_names={},
polyaxon_init=V1PolyaxonInitContainer(image="foo/foo"),
)
assert containers == []
# With context
containers = self.converter.get_init_containers(
contexts=PluginsContextsSpec.from_config(
V1Plugins(collect_artifacts=True,
collect_logs=False,
auth=True)),
artifacts_store=store,
init_connections=[],
init_containers=[],
connection_by_names={},
polyaxon_init=V1PolyaxonInitContainer(image="foo/foo"),
)
assert containers == [
get_auth_context_container(
polyaxon_init=V1PolyaxonInitContainer(image="foo/foo"),
env=self.converter.get_auth_service_env_vars(),
),
get_artifacts_path_container(
polyaxon_init=V1PolyaxonInitContainer(image="foo/foo"),
artifacts_store=store,
run_path=self.converter.run_path,
auto_resume=True,
),
]
def get_resource(
self,
compiled_operation: V1CompiledOperation,
artifacts_store: V1ConnectionType,
connection_by_names: Dict[str, V1ConnectionType],
secrets: Optional[Iterable[V1K8sResourceType]],
config_maps: Optional[Iterable[V1K8sResourceType]],
is_resume: bool = False,
) -> Dict:
job = compiled_operation.run # type: V1Job
plugins = compiled_operation.plugins or V1Plugins()
contexts = PluginsContextsSpec.from_config(plugins)
replica_spec = self.get_replica_resource(
plugins=plugins,
contexts=contexts,
environment=job.environment,
volumes=job.volumes or [],
init=job.init or [],
sidecars=job.sidecars or [],
container=job.container,
artifacts_store=artifacts_store,
connections=job.connections or [],
connection_by_names=connection_by_names,
secrets=secrets,
config_maps=config_maps,
is_resume=is_resume,
)
return get_job_custom_resource(
namespace=self.namespace,
main_container=replica_spec.main_container,
sidecar_containers=replica_spec.sidecar_containers,
init_containers=replica_spec.init_containers,
resource_name=self.resource_name,
volumes=replica_spec.volumes,
environment=replica_spec.environment,
termination=compiled_operation.termination,
collect_logs=contexts.collect_logs,
sync_statuses=contexts.sync_statuses,
notifications=plugins.notifications,
labels=replica_spec.labels,
)
def get_cleaner_operation(connection: V1ConnectionType, run_uuid: str,
run_kind: str) -> V1Operation:
return V1Operation(
termination=V1Termination(max_retries=1),
component=V1Component(
name="cleaner",
plugins=V1Plugins(
auth=False,
collect_logs=False,
collect_artifacts=False,
collect_resources=False,
auto_resume=False,
sync_statuses=False,
),
run=V1CleanerJob(
connections=[connection.name],
container=get_default_cleaner_container(
connection, run_uuid, run_kind),
),
),
)
def get_batch_cleaner_operation(
connection: V1ConnectionType,
paths: List[str],
) -> V1Operation:
return V1Operation(
termination=V1Termination(max_retries=1),
component=V1Component(
name="cleaner",
plugins=V1Plugins(
auth=False,
collect_logs=False,
collect_artifacts=False,
collect_resources=False,
auto_resume=False,
sync_statuses=False,
),
run=V1CleanerJob(
connections=[connection.name],
container=get_batch_cleaner_container(connection, paths),
),
),
)
def test_all_volumes(self):
assert (
len(
get_volume_mounts(
contexts=PluginsContextsSpec.from_config(
V1Plugins(collect_logs=False, collect_artifacts=True)),
init=[
V1Init(connection=self.s3_store.name, path="/test-1"),
V1Init(connection=self.gcs_store.name, path="/test-2"),
V1Init(connection=self.az_store.name, path="/test-3"),
V1Init(connection=self.claim_store.name,
path="/test-4"),
V1Init(connection=self.host_path_store.name,
path="/test-5"),
],
connections=[
self.s3_store,
self.gcs_store,
self.az_store,
self.claim_store,
self.host_path_store,
],
secrets=[
self.non_mount_resource1,
self.non_mount_resource1,
self.mount_resource1,
self.mount_resource2,
],
config_maps=[
self.non_mount_resource1,
self.non_mount_resource1,
self.mount_resource1,
self.mount_resource2,
],
))
# 1: output store
# 7: 5 managed contexts + 2 mounts
# 4: 4 mount resources (secrets + configs)
== 1 + 7 + 4)
assert (
len(
get_volume_mounts(
contexts=None,
init=[
V1Init(connection=self.s3_store.name, path="/test-1"),
V1Init(connection=self.gcs_store.name, path="/test-2"),
V1Init(connection=self.az_store.name, path="/test-3"),
V1Init(connection=self.claim_store.name,
path="/test-4"),
V1Init(connection=self.host_path_store.name,
path="/test-5"),
],
connections=[
self.s3_store,
self.gcs_store,
self.az_store,
self.claim_store,
self.host_path_store,
],
secrets=[
self.non_mount_resource1,
self.non_mount_resource1,
self.mount_resource1,
self.mount_resource2,
],
config_maps=[
self.non_mount_resource1,
self.non_mount_resource1,
self.mount_resource1,
self.mount_resource2,
],
))
# 7: 5 managed contexts + 2 mounts
# 4: 4 mount resources (secrets + configs)
== 7 + 4)
assert (
len(
get_volume_mounts(
contexts=PluginsContextsSpec.from_config(
V1Plugins(collect_logs=True, collect_artifacts=True)),
init=[
V1Init(connection=self.s3_store.name, path="/test-1"),
V1Init(connection=self.gcs_store.name, path="/test-2"),
V1Init(connection=self.az_store.name, path="/test-3"),
V1Init(connection=self.claim_store.name,
path="/test-4"),
V1Init(connection=self.host_path_store.name,
path="/test-5"),
],
connections=[
self.s3_store,
self.gcs_store,
self.az_store,
self.claim_store,
self.host_path_store,
],
secrets=[
self.non_mount_resource1,
self.non_mount_resource1,
self.mount_resource1,
self.mount_resource2,
],
config_maps=[
self.non_mount_resource1,
self.non_mount_resource1,
self.mount_resource1,
self.mount_resource2,
],
))
# 1: outputs context store
# 7: 5 managed contexts + 2 mounts
# 4: 4 mount resources (secrets + configs)
== 1 + 7 + 4)
assert (
len(
get_volume_mounts(
contexts=PluginsContextsSpec.from_config(
V1Plugins(collect_logs=True, collect_artifacts=False)),
init=[
V1Init(connection=self.s3_store.name, path="/test-1"),
V1Init(connection=self.gcs_store.name, path="/test-2"),
V1Init(connection=self.az_store.name, path="/test-3"),
V1Init(connection=self.claim_store.name,
path="/test-4"),
V1Init(connection=self.host_path_store.name,
path="/test-5"),
],
connections=[
self.s3_store,
self.gcs_store,
self.az_store,
self.claim_store,
self.host_path_store,
],
secrets=[
self.non_mount_resource1,
self.non_mount_resource1,
self.mount_resource1,
self.mount_resource2,
],
config_maps=[
self.non_mount_resource1,
self.non_mount_resource1,
self.mount_resource1,
self.mount_resource2,
],
))
# 7: 5 managed contexts + 2 mounts
# 4: 4 mount resources (secrets + configs)
== 7 + 4)
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
