Skip to content

Commit

Permalink
Revert "Split create_or_find_service_account into a non-mutating meth…
Browse files Browse the repository at this point in the history
…od and a mutating method"
  • Loading branch information
ankit864 authored Jul 2, 2024
1 parent 10bdad1 commit 8397ccd
Show file tree
Hide file tree
Showing 6 changed files with 70 additions and 59 deletions.
20 changes: 0 additions & 20 deletions paasta_tools/kubernetes/application/controller_wrappers.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
from paasta_tools.kubernetes_tools import create_deployment
from paasta_tools.kubernetes_tools import create_pod_disruption_budget
from paasta_tools.kubernetes_tools import create_stateful_set
from paasta_tools.kubernetes_tools import ensure_service_account
from paasta_tools.kubernetes_tools import KubeClient
from paasta_tools.kubernetes_tools import KubeDeployment
from paasta_tools.kubernetes_tools import KubernetesDeploymentConfig
Expand Down Expand Up @@ -121,25 +120,6 @@ def update_related_api_objects(self, kube_client: KubeClient) -> None:
"""
self.ensure_pod_disruption_budget(kube_client, self.soa_config.get_namespace())

def update_dependency_api_objects(self, kube_client: KubeClient) -> None:
"""
Update related Kubernetes API objects that should be updated before the main object,
such as service accounts.
:param kube_client:
"""
self.ensure_service_account(kube_client)

def ensure_service_account(self, kube_client: KubeClient) -> None:
"""
Ensure that the service account for this application exists
:param kube_client:
"""
ensure_service_account(
iam_role=self.soa_config.get_iam_role(),
namespace=self.soa_config.get_namespace(),
kube_client=kube_client,
)

def delete_pod_disruption_budget(self, kube_client: KubeClient) -> None:
try:
kube_client.policy.delete_namespaced_pod_disruption_budget(
Expand Down
28 changes: 14 additions & 14 deletions paasta_tools/kubernetes_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -2220,9 +2220,9 @@ def get_pod_template_spec(
annotations["iam.amazonaws.com/role"] = ""
iam_role = self.get_iam_role()
if iam_role:
pod_spec_kwargs["service_account_name"] = get_service_account_name(
iam_role
)
pod_spec_kwargs[
"service_account_name"
] = create_or_find_service_account_name(iam_role, self.get_namespace())
if fs_group is None:
# We need some reasoable default for group id of a process
# running inside the container. Seems like most of such
Expand Down Expand Up @@ -4050,9 +4050,12 @@ def get_all_limit_ranges(
_RE_NORMALIZE_IAM_ROLE = re.compile(r"[^0-9a-zA-Z]+")


def get_service_account_name(
def create_or_find_service_account_name(
iam_role: str,
namespace: str,
k8s_role: Optional[str] = None,
kubeconfig_file: Optional[str] = None,
dry_run: bool = False,
) -> str:
# the service account is expected to always be prefixed with paasta- as using the actual namespace
# potentially wastes a lot of characters (e.g., paasta-nrtsearchservices) that could be used for
Expand All @@ -4078,17 +4081,12 @@ def get_service_account_name(
"Expected at least one of iam_role or k8s_role to be passed in!"
)

return sa_name


def ensure_service_account(
iam_role: str,
namespace: str,
kube_client: KubeClient,
k8s_role: Optional[str] = None,
) -> None:
sa_name = get_service_account_name(iam_role, k8s_role)
# if someone is dry-running paasta_setup_tron_namespace or some other tool that
# calls this function, we probably don't want to mutate k8s state :)
if dry_run:
return sa_name

kube_client = KubeClient(config_file=kubeconfig_file)
if not any(
sa.metadata and sa.metadata.name == sa_name
for sa in get_all_service_accounts(kube_client, namespace)
Expand Down Expand Up @@ -4137,6 +4135,8 @@ def ensure_service_account(
namespace=namespace, body=role_binding
)

return sa_name


def mode_to_int(mode: Optional[Union[str, int]]) -> Optional[int]:
if mode is not None:
Expand Down
1 change: 0 additions & 1 deletion paasta_tools/setup_kubernetes_job.py
Original file line number Diff line number Diff line change
Expand Up @@ -281,7 +281,6 @@ def setup_kube_deployments(
"paasta_namespace": app.kube_deployment.namespace,
}
try:
app.update_dependency_api_objects(kube_client)
if (
app.kube_deployment.service,
app.kube_deployment.instance,
Expand Down
16 changes: 13 additions & 3 deletions paasta_tools/tron_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,7 @@

from paasta_tools.kubernetes_tools import (
allowlist_denylist_to_requirements,
get_service_account_name,
create_or_find_service_account_name,
limit_size_with_hash,
raw_selectors_to_requirements,
to_node_label,
Expand Down Expand Up @@ -354,11 +354,15 @@ def build_spark_config(self) -> Dict[str, str]:
"spark.kubernetes.executor.label.yelp.com/owner", self.get_team()
)

# We need to make sure the Service Account used by the executors has been created.
# We are using the Service Account created using the provided or default IAM role.
spark_conf[
"spark.kubernetes.authenticate.executor.serviceAccountName"
] = get_service_account_name(
] = create_or_find_service_account_name(
iam_role=self.get_spark_executor_iam_role(),
namespace=spark_tools.SPARK_EXECUTOR_NAMESPACE,
kubeconfig_file=system_paasta_config.get_spark_kubeconfig(),
dry_run=self.for_validation,
)

return spark_conf
Expand Down Expand Up @@ -948,14 +952,20 @@ def format_tron_action_dict(action_config: TronActionConfig):

result["labels"]["yelp.com/owner"] = "compute_infra_platform_experience"

# create_or_find_service_account_name requires k8s credentials, and we don't
# have those available for CI to use (nor do we check these for normal PaaSTA
# services, so we're not doing anything "new" by skipping this)
if (
action_config.get_iam_role_provider() == "aws"
and action_config.get_iam_role()
and not action_config.for_validation
):
# this service account will be used for normal Tron batches as well as for Spark drivers
result["service_account_name"] = get_service_account_name(
result["service_account_name"] = create_or_find_service_account_name(
iam_role=action_config.get_iam_role(),
namespace=EXECUTOR_TYPE_TO_NAMESPACE[executor],
k8s_role=None,
dry_run=action_config.for_validation,
)

# service account token volumes for service authentication
Expand Down
62 changes: 42 additions & 20 deletions tests/test_kubernetes_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -89,14 +89,14 @@
from paasta_tools.kubernetes_tools import allowlist_denylist_to_requirements
from paasta_tools.kubernetes_tools import create_custom_resource
from paasta_tools.kubernetes_tools import create_deployment
from paasta_tools.kubernetes_tools import create_or_find_service_account_name
from paasta_tools.kubernetes_tools import create_pod_disruption_budget
from paasta_tools.kubernetes_tools import create_secret
from paasta_tools.kubernetes_tools import create_secret_signature
from paasta_tools.kubernetes_tools import create_stateful_set
from paasta_tools.kubernetes_tools import ensure_namespace
from paasta_tools.kubernetes_tools import ensure_paasta_api_rolebinding
from paasta_tools.kubernetes_tools import ensure_paasta_namespace_limits
from paasta_tools.kubernetes_tools import ensure_service_account
from paasta_tools.kubernetes_tools import filter_nodes_by_blacklist
from paasta_tools.kubernetes_tools import filter_pods_by_service_instance
from paasta_tools.kubernetes_tools import force_delete_pods
Expand All @@ -117,7 +117,6 @@
from paasta_tools.kubernetes_tools import get_secret
from paasta_tools.kubernetes_tools import get_secret_name_from_ref
from paasta_tools.kubernetes_tools import get_secret_signature
from paasta_tools.kubernetes_tools import get_service_account_name
from paasta_tools.kubernetes_tools import InvalidKubernetesConfig
from paasta_tools.kubernetes_tools import is_node_ready
from paasta_tools.kubernetes_tools import is_pod_ready
Expand Down Expand Up @@ -4458,7 +4457,7 @@ def test_get_pod_management_policy(config_dict, expected_management_policy):
assert deployment.get_pod_management_policy() == expected_management_policy


def test_ensure_service_account_new():
def test_create_or_find_service_account_name_new():
iam_role = "arn:aws:iam::000000000000:role/some_role"
namespace = "test_namespace"
k8s_role = None
Expand All @@ -4478,8 +4477,8 @@ def test_ensure_service_account_new():
mock_client.core.list_namespaced_service_account.return_value.items = []
mock_kube_client.return_value = mock_client

ensure_service_account(
iam_role, namespace=namespace, kube_client=mock_client, k8s_role=k8s_role
assert expected_sa_name == create_or_find_service_account_name(
iam_role, namespace=namespace, k8s_role=k8s_role
)
mock_client.core.create_namespaced_service_account.assert_called_once_with(
namespace=namespace,
Expand All @@ -4495,7 +4494,7 @@ def test_ensure_service_account_new():
mock_client.rbac.create_namespaced_role_binding.assert_not_called()


def test_ensure_service_account_with_k8s_role_new():
def test_create_or_find_service_account_name_with_k8s_role_new():
iam_role = "arn:aws:iam::000000000000:role/some_role"
namespace = "test_namespace"
k8s_role = "mega-admin"
Expand All @@ -4519,8 +4518,8 @@ def test_ensure_service_account_with_k8s_role_new():
mock_client.rbac.list_namespaced_role_binding.return_value.items = []
mock_kube_client.return_value = mock_client

ensure_service_account(
iam_role, namespace=namespace, kube_client=mock_client, k8s_role=k8s_role
assert expected_sa_name == create_or_find_service_account_name(
iam_role, namespace=namespace, k8s_role=k8s_role
)
mock_client.core.create_namespaced_service_account.assert_called_once_with(
namespace=namespace,
Expand Down Expand Up @@ -4556,7 +4555,7 @@ def test_ensure_service_account_with_k8s_role_new():
)


def test_ensure_service_account_existing():
def test_create_or_find_service_account_name_existing():
iam_role = "arn:aws:iam::000000000000:role/some_role"
namespace = "test_namespace"
k8s_role = "mega-admin"
Expand Down Expand Up @@ -4609,14 +4608,14 @@ def test_ensure_service_account_existing():
]
mock_kube_client.return_value = mock_client

ensure_service_account(
iam_role, namespace=namespace, kube_client=mock_client, k8s_role=k8s_role
assert expected_sa_name == create_or_find_service_account_name(
iam_role, namespace=namespace, k8s_role=k8s_role
)
mock_client.core.create_namespaced_service_account.assert_not_called()
mock_client.rbac.create_namespaced_role_binding.assert_not_called()


def test_ensure_service_account_existing_create_rb_only():
def test_create_or_find_service_account_name_existing_create_rb_only():
iam_role = "arn:aws:iam::000000000000:role/some_role"
namespace = "test_namespace"
k8s_role = "mega-admin"
Expand Down Expand Up @@ -4649,25 +4648,48 @@ def test_ensure_service_account_existing_create_rb_only():
mock_client.rbac.list_namespaced_role_binding.return_value.items = []
mock_kube_client.return_value = mock_client

ensure_service_account(
iam_role, namespace=namespace, kube_client=mock_client, k8s_role=k8s_role
assert expected_sa_name == create_or_find_service_account_name(
iam_role, namespace=namespace, k8s_role=k8s_role
)
mock_client.core.create_namespaced_service_account.assert_not_called()
assert mock_client.rbac.create_namespaced_role_binding.called is True


def test_ensure_service_account_caps():
def test_create_or_find_service_account_name_caps():
iam_role = "arn:aws:iam::000000000000:role/Some_Role"
namespace = "test_namespace"
expected_sa_name = "paasta--arn-aws-iam-000000000000-role-some-role"
with mock.patch(
"paasta_tools.kubernetes_tools.kube_config.load_kube_config", autospec=True
):
assert expected_sa_name == get_service_account_name(
), mock.patch(
"paasta_tools.kubernetes_tools.KubeClient",
autospec=False,
) as mock_kube_client:
mock_client = mock.Mock()
mock_client.core = mock.Mock(spec=kube_client.CoreV1Api)
mock_client.core.list_namespaced_service_account.return_value = mock.Mock(
spec=V1ServiceAccountList
)
mock_client.core.list_namespaced_service_account.return_value.items = [
V1ServiceAccount(
kind="ServiceAccount",
metadata=V1ObjectMeta(
name=expected_sa_name,
namespace=namespace,
annotations={"eks.amazonaws.com/role-arn": iam_role},
),
)
]
mock_kube_client.return_value = mock_client

assert expected_sa_name == create_or_find_service_account_name(
iam_role,
namespace=namespace,
)
mock_client.core.create_namespaced_service_account.assert_not_called()


def test_ensure_service_account_caps_with_k8s():
def test_create_or_find_service_account_name_caps_with_k8s():
iam_role = "arn:aws:iam::000000000000:role/Some_Role"
namespace = "test_namespace"
k8s_role = "mega-admin"
Expand Down Expand Up @@ -4720,8 +4742,8 @@ def test_ensure_service_account_caps_with_k8s():
]
mock_kube_client.return_value = mock_client

ensure_service_account(
iam_role, namespace=namespace, kube_client=mock_client, k8s_role=k8s_role
assert expected_sa_name == create_or_find_service_account_name(
iam_role, namespace=namespace, k8s_role=k8s_role
)
mock_client.core.create_namespaced_service_account.assert_not_called()
mock_client.rbac.create_namespaced_role_binding.assert_not_called()
Expand Down
2 changes: 1 addition & 1 deletion tests/test_tron_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -1426,7 +1426,7 @@ def test_format_tron_action_dict_paasta_k8s(
autospec=True,
return_value=False,
), mock.patch(
"paasta_tools.tron_tools.get_service_account_name",
"paasta_tools.tron_tools.create_or_find_service_account_name",
autospec=True,
return_value="some--service--account",
), mock.patch(
Expand Down

0 comments on commit 8397ccd

Please sign in to comment.