def test_agent_config(self):
config_dict = {POLYAXON_KEYS_AGENT_ARTIFACTS_STORE: 12}
with self.assertRaises(ValidationError):
AgentConfig.from_dict(config_dict)
config_dict = {POLYAXON_KEYS_AGENT_ARTIFACTS_STORE: "some"}
with self.assertRaises(ValidationError):
AgentConfig.from_dict(config_dict)
config_dict = {
POLYAXON_KEYS_K8S_NAMESPACE: "foo",
POLYAXON_KEYS_AGENT_ARTIFACTS_STORE: {
"name": "some",
"kind": V1ConnectionKind.GCS,
"schema": V1BucketConnection(bucket="gs://test").to_dict(),
},
POLYAXON_KEYS_AGENT_CONNECTIONS: [],
}
config = AgentConfig.from_dict(config_dict)
assert config.to_light_dict() == config_dict
config_dict = {
POLYAXON_KEYS_K8S_NAMESPACE: "foo",
POLYAXON_KEYS_AGENT_ARTIFACTS_STORE: "some",
POLYAXON_KEYS_AGENT_CONNECTIONS: [
{
"name": "some",
"kind": V1ConnectionKind.GCS,
"schema": V1BucketConnection(bucket="gs://test").to_dict(),
"secretResource": "some",
}
],
}
with self.assertRaises(ValidationError):
AgentConfig.from_dict(config_dict)
config_dict = {
POLYAXON_KEYS_K8S_NAMESPACE: "foo",
POLYAXON_KEYS_AGENT_ARTIFACTS_STORE: {
"name": "test",
"kind": V1ConnectionKind.GCS,
"schema": V1BucketConnection(bucket="gs://test").to_dict(),
"secret": V1K8sResourceSchema(name="some").to_dict(),
},
POLYAXON_KEYS_AGENT_CONNECTIONS: [
{
"name": "some",
"kind": V1ConnectionKind.GCS,
"schema": V1BucketConnection(bucket="gs://test").to_dict(),
"secret": V1K8sResourceSchema(name="some").to_dict(),
},
{
"name": "slack",
"kind": V1ConnectionKind.SLACK,
"secret": V1K8sResourceSchema(name="some").to_dict(),
},
],
}
config = AgentConfig.from_dict(config_dict)
assert config.to_light_dict() == config_dict
def convert(
owner_name: str,
project_name: str,
run_name: str,
run_uuid: str,
content: str,
default_auth: bool,
agent_content: str = None,
) -> Dict:
polypod_config = PolypodConfig()
compiled_operation = CompiledOperationSpecification.read(content)
polypod_config.resolve(
compiled_operation=compiled_operation,
agent_config=AgentConfig.read(agent_content)
if agent_content else None,
)
return converter.convert(
compiled_operation=compiled_operation,
owner_name=owner_name,
project_name=project_name,
run_name=run_name,
run_uuid=run_uuid,
namespace=polypod_config.namespace,
polyaxon_init=polypod_config.polyaxon_init,
polyaxon_sidecar=polypod_config.polyaxon_sidecar,
run_path=run_uuid,
artifacts_store=polypod_config.artifacts_store,
connection_by_names=polypod_config.connection_by_names,
secrets=polypod_config.secrets,
config_maps=polypod_config.config_maps,
default_sa=polypod_config.default_sa,
converters=PLATFORM_CONVERTERS,
default_auth=default_auth,
)
def test_agent_config_from_str_envs(self):
config_dict = {
POLYAXON_KEYS_K8S_NAMESPACE: "foo",
POLYAXON_KEYS_AGENT_ARTIFACTS_STORE: ujson.dumps(
{
"name": "test1",
"kind": V1ConnectionKind.GCS,
"schema": V1BucketConnection(bucket="gs://test").to_dict(),
"secret": V1K8sResourceSchema(name="some").to_dict(),
}
),
POLYAXON_KEYS_AGENT_NOTIFICATION_CONNECTIONS: ujson.dumps(
[
{
"name": "slack",
"kind": V1ConnectionKind.SLACK,
"secret": V1K8sResourceSchema(name="some").to_dict(),
}
],
),
POLYAXON_KEYS_AGENT_CONNECTIONS: ujson.dumps(
[
{
"name": "test2",
"kind": V1ConnectionKind.GCS,
"schema": V1BucketConnection(bucket="gs://test").to_dict(),
"secret": V1K8sResourceSchema(name="some").to_dict(),
}
]
),
}
config = AgentConfig.from_dict(config_dict)
assert len(config.secrets) == 1
assert len(config.to_light_dict()[POLYAXON_KEYS_AGENT_CONNECTIONS]) == 1
def patch_settings(
set_auth: bool = True,
set_client: bool = True,
set_cli: bool = True,
set_agent: bool = True,
set_proxies: bool = True,
):
settings.AUTH_CONFIG = None
if set_auth:
settings.AUTH_CONFIG = AccessTokenConfig()
settings.CLIENT_CONFIG = None
if set_client:
settings.CLIENT_CONFIG = ClientConfig(host="1.2.3.4")
settings.CLI_CONFIG = None
if set_cli:
settings.CLI_CONFIG = CliConfig()
settings.AGENT_CONFIG = None
if set_agent:
settings.AGENT_CONFIG = AgentConfig()
settings.PROXIES_CONFIG = None
if set_proxies:
settings.PROXIES_CONFIG = ProxiesConfig()
settings.CLIENT_CONFIG.tracking_timeout = 0
def test_resolve_connections_with_invalid_config(self):
fpath = tempfile.mkdtemp()
AgentConfigManager.CONFIG_PATH = fpath
secret1 = V1K8sResourceType(
name="secret1",
schema=V1K8sResourceSchema(name="secret1"),
is_requested=True,
)
secret2 = V1K8sResourceType(
name="secret2",
schema=V1K8sResourceSchema(name="secret2"),
is_requested=True,
)
artifacts_store = V1ConnectionType(
name="test_s3",
kind=V1ConnectionKind.S3,
schema=V1BucketConnection(bucket="s3//:foo"),
secret=secret1.schema,
)
connection1 = V1ConnectionType(
name="connection1",
kind=V1ConnectionKind.REGISTRY,
schema=V1HostConnection(url="localhost:5000"),
secret=secret2.schema,
)
connection2 = V1ConnectionType(
name="connection2",
kind=V1ConnectionKind.REGISTRY,
)
settings.AGENT_CONFIG = AgentConfig(
namespace="foo",
artifacts_store=artifacts_store,
connections=[connection1, connection2],
)
resolver = BaseResolver(
run=None,
compiled_operation=self.compiled_operation,
owner_name="user",
project_name="p1",
project_uuid=None,
run_name="j1",
run_uuid=None,
run_path="test",
params=None,
)
resolver.resolve_connections()
assert resolver.namespace == "foo"
assert resolver.connection_by_names == {
artifacts_store.name: artifacts_store,
connection1.name: connection1,
connection2.name: connection2,
}
assert resolver.artifacts_store == artifacts_store
assert resolver.polyaxon_sidecar == get_default_sidecar_container()
assert resolver.polyaxon_init == get_default_init_container()
resolver.resolve_artifacts_lineage()
assert len(resolver.artifacts) == 3
def set_store():
store_root = tempfile.mkdtemp()
settings.AGENT_CONFIG = AgentConfig(
artifacts_store=V1ConnectionType(
name="test",
kind=V1ConnectionKind.HOST_PATH,
schema=V1HostPathConnection(host_path=store_root,
mount_path=store_root),
secret=None,
),
connections=[],
)
settings.CLIENT_CONFIG.archive_root = tempfile.mkdtemp()
return store_root
def get_config_from_env(cls) -> AgentConfig:
tmp_path = os.path.join(CONTEXT_TMP_POLYAXON_PATH, ".agent")
user_path = os.path.join(CONTEXT_USER_POLYAXON_PATH, ".agent")
config_paths = [
os.environ,
ConfigSpec(tmp_path, config_type=".json", check_if_exists=False),
ConfigSpec(user_path, config_type=".json", check_if_exists=False),
{
"dummy": "dummy"
},
]
agent_config = ConfigManager.read_configs(config_paths)
return AgentConfig.from_dict(agent_config.data)
def setUp(self):
super().setUp()
settings.AUTH_CONFIG = None
if self.SET_AUTH_SETTINGS:
settings.AUTH_CONFIG = AccessTokenConfig()
settings.CLIENT_CONFIG = None
if self.SET_CLIENT_SETTINGS:
settings.CLIENT_CONFIG = ClientConfig()
settings.AGENT_CONFIG = None
if self.SET_AGENT_SETTINGS:
settings.AGENT_CONFIG = AgentConfig()
settings.PROXIES_CONFIG = None
if self.SET_PROXIES_SETTINGS:
settings.PROXIES_CONFIG = ProxiesConfig()
fpath = tempfile.mkdtemp()
settings.CLIENT_CONFIG.agent_path = fpath
def test_resolve_connections_with_invalid_config(self):
fpath = tempfile.mkdtemp()
AgentManager.CONFIG_PATH = fpath
secret1 = V1K8sResourceType(
name="secret1",
schema=V1K8sResourceSchema(name="secret1"),
is_requested=True,
)
secret2 = V1K8sResourceType(
name="secret2",
schema=V1K8sResourceSchema(name="secret2"),
is_requested=True,
)
connection1 = V1ConnectionType(
name="test_s3",
kind=V1ConnectionKind.S3,
schema=V1BucketConnection(bucket="s3//:foo"),
secret=secret1.schema,
)
connection2 = V1ConnectionType(
name="test_gcs",
kind=V1ConnectionKind.GCS,
schema=V1BucketConnection(bucket="gcs//:foo"),
secret=secret1.schema,
)
connection3 = V1ConnectionType(
name="test_wasb",
kind=V1ConnectionKind.WASB,
schema=V1BucketConnection(bucket="wasbs//:foo"),
secret=secret2.schema,
)
settings.AGENT_CONFIG = AgentConfig(
namespace="foo",
artifacts_store=connection1,
connections=[connection2, connection3],
)
resolver = BaseResolver(
run=None,
compiled_operation=self.compiled_operation,
owner_name="user",
project_name="p1",
project_uuid=None,
run_name="j1",
run_uuid=None,
run_path="test",
params=None,
)
resolver.resolve_connections()
assert resolver.namespace == "foo"
assert resolver.connection_by_names == {connection1.name: connection1}
assert resolver.artifacts_store == connection1
assert [s.schema
for s in resolver.secrets] == [secret1.schema, secret2.schema]
assert resolver.polyaxon_sidecar == get_default_sidecar_container()
assert resolver.polyaxon_init == get_default_init_container()
# Add run spec to resolve connections
compiled_operation = V1CompiledOperation.read({
"version": 1.1,
"kind": kinds.COMPILED_OPERATION,
"plugins": {
"auth": False,
"shm": False,
"collectLogs": False,
"collectArtifacts": False,
"collectResources": False,
},
"run": {
"kind": V1RunKind.JOB,
"container": {
"image": "test"
},
"connections": {connection3.name},
},
})
resolver = BaseResolver(
run=None,
compiled_operation=compiled_operation,
owner_name="user",
project_name="p1",
project_uuid=None,
run_name="j1",
run_uuid=None,
run_path="test",
params=None,
)
resolver.resolve_connections()
assert resolver.namespace == "foo"
assert resolver.connection_by_names == {
connection1.name: connection1,
connection3.name: connection3,
}
assert [s.schema
for s in resolver.secrets] == [secret1.schema, secret2.schema]
assert resolver.artifacts_store == connection1
assert resolver.polyaxon_sidecar == get_default_sidecar_container()
assert resolver.polyaxon_init == get_default_init_container()
# Add run spec to resolve connections
compiled_operation = V1CompiledOperation.read({
"version": 1.1,
"kind": kinds.COMPILED_OPERATION,
"plugins": {
"auth": False,
"shm": False,
"collectLogs": False,
"collectArtifacts": False,
"collectResources": False,
},
"run": {
"kind": V1RunKind.JOB,
"container": {
"image": "test"
},
"connections": {
connection1.name,
connection2.name,
connection3.name,
},
},
})
resolver = BaseResolver(
run=None,
compiled_operation=compiled_operation,
owner_name="user",
project_name="p1",
project_uuid=None,
run_name="j1",
run_uuid=None,
run_path="test",
params=None,
)
resolver.resolve_connections()
assert resolver.namespace == "foo"
assert resolver.connection_by_names == {
connection3.name: connection3,
connection2.name: connection2,
connection1.name: connection1,
}
assert [s.schema
for s in resolver.secrets] == [secret1.schema, secret2.schema]
assert resolver.artifacts_store == connection1
assert resolver.polyaxon_sidecar == get_default_sidecar_container()
assert resolver.polyaxon_init == get_default_init_container()