Skip to content

Commit

Permalink
refactor: refactor env variables for kfp-api to match v2.0.2 (#354)
Browse files Browse the repository at this point in the history 
* refactor: refactor env variables for kfp-api to match v2.0.2

For a few releases the KFP API server component has been defining
certain configurations using env variables instead of a config file.
This change plus adding/removing vars will ensure the charmed kfp-api
keeps consistent with the upstream project.
  • Loading branch information
@DnPlas
DnPlas authored Oct 26, 2023
1 parent 3bb5478 commit 468b6e7
Show file tree
Hide file tree
Showing 3 changed files with 100 additions and 129 deletions.
130 changes: 40 additions & 90 deletions charms/kfp-api/src/charm.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
https://github.com/canonical/kfp-operators/
"""

import json
import logging
from pathlib import Path

Expand All @@ -24,14 +23,7 @@
from lightkube.models.core_v1 import ServicePort
from ops.charm import CharmBase
from ops.main import main
from ops.model import (
ActiveStatus,
BlockedStatus,
Container,
MaintenanceStatus,
ModelError,
WaitingStatus,
)
from ops.model import ActiveStatus, BlockedStatus, MaintenanceStatus, ModelError, WaitingStatus
from ops.pebble import CheckStatus, Layer
from serialized_data_interface import (
NoCompatibleVersions,
Expand Down Expand Up @@ -175,9 +167,7 @@ def k8s_resource_handler(self, handler: KubernetesResourceHandler):
@property
def service_environment(self):
"""Return environment variables based on model configuration."""
ret_env_vars = {"POD_NAMESPACE": self.model.name}

return ret_env_vars
return self._generate_environment()

@property
def _kfp_api_layer(self) -> Layer:
Expand Down Expand Up @@ -210,76 +200,61 @@ def _kfp_api_layer(self) -> Layer:

return Layer(layer_config)

def _generate_config(self, interfaces):
"""Generate configuration based on supplied data.
def _generate_environment(self) -> dict:
"""Generate environment based on supplied data.
Configuration is generated based on:
- Supplied interfaces.
- Database data: from MySQL relation data or from data platform library.
- Model configuration.
Return:
env_vars(dict): a dictionary of environment variables for the api server.
"""

config = self.model.config
try:
interfaces = self._get_interfaces()
db_data = self._get_db_data()
os = self._get_object_storage(interfaces)
viz = self._get_viz(interfaces)
object_storage = self._get_object_storage(interfaces)
viz_data = self._get_viz(interfaces)
except ErrorWithStatus as error:
self.logger.error("Failed to generate container configuration.")
raise error

# at this point all data is correctly populated and proper config can be generated
config_json = {
"DBConfig": {
"ConMaxLifeTime": "120s",
"DBName": db_data["db_name"],
"DriverName": "mysql",
"GroupConcatMaxLen": "4194304",
"Host": db_data["db_host"],
"Password": db_data["db_password"],
"Port": db_data["db_port"],
"User": db_data["db_username"],
},
"ObjectStoreConfig": {
"AccessKey": os["access-key"],
"BucketName": config["object-store-bucket-name"],
"Host": f"{os['service']}.{os['namespace']}",
"Multipart": {"Disable": "true"},
"PipelinePath": "pipelines",
"Port": str(os["port"]),
"Region": "",
"SecretAccessKey": os["secret-key"],
"Secure": str(os["secure"]).lower(),
},
"ARCHIVE_CONFIG_LOG_FILE_NAME": config["log-archive-filename"],
"ARCHIVE_CONFIG_LOG_PATH_PREFIX": config["log-archive-prefix"],
"AUTO_UPDATE_PIPELINE_DEFAULT_VERSION": str(
config["auto-update-default-version"]
).lower(),
"CACHE_IMAGE": config["cache-image"],
"CACHE_NODE_RESTRICTIONS": "false",
"CacheEnabled": str(config["cache-enabled"]).lower(),
"DefaultPipelineRunnerServiceAccount": config["runner-sa"],
"InitConnectionTimeout": config["init-connection-timeout"],
env_vars = {
"AUTO_UPDATE_PIPELINE_DEFAULT_VERSION": self.model.config[
"auto-update-default-version"
],
"KFP_API_SERVICE_NAME": KFP_API_SERVICE_NAME,
"KUBEFLOW_USERID_HEADER": "kubeflow-userid",
"KUBEFLOW_USERID_PREFIX": "",
"POD_NAMESPACE": self.model.name,
"OBJECTSTORECONFIG_SECURE": "false",
"OBJECTSTORECONFIG_BUCKETNAME": self.model.config["object-store-bucket-name"],
"DBCONFIG_USER": db_data["db_username"],
"DBCONFIG_PASSWORD": db_data["db_password"],
"DBCONFIG_DBNAME": db_data["db_name"],
"DBCONFIG_HOST": db_data["db_host"],
"DBCONFIG_PORT": db_data["db_port"],
"DBCONFIG_CONMAXLIFETIME": "120s",
"DB_DRIVER_NAME": "mysql",
"DBCONFIG_MYSQLCONFIG_USER": db_data["db_username"],
"DBCONFIG_MYSQLCONFIG_PASSWORD": db_data["db_password"],
"DBCONFIG_MYSQLCONFIG_DBNAME": db_data["db_name"],
"DBCONFIG_MYSQLCONFIG_HOST": db_data["db_host"],
"DBCONFIG_MYSQLCONFIG_PORT": db_data["db_port"],
"OBJECTSTORECONFIG_ACCESSKEY": object_storage["access-key"],
"OBJECTSTORECONFIG_SECRETACCESSKEY": object_storage["secret-key"],
"DEFAULTPIPELINERUNNERSERVICEACCOUNT": "default-editor",
"MULTIUSER": "true",
"ML_PIPELINE_VISUALIZATIONSERVER_SERVICE_HOST": viz["service-name"],
"ML_PIPELINE_VISUALIZATIONSERVER_SERVICE_PORT": viz["service-port"],
"VISUALIZATIONSERVICE_NAME": viz_data["service-name"],
"VISUALIZATIONSERVICE_PORT": viz_data["service-port"],
"ML_PIPELINE_VISUALIZATIONSERVER_SERVICE_HOST": viz_data["service-name"],
"ML_PIPELINE_VISUALIZATIONSERVER_SERVICE_PORT": viz_data["service-port"],
"CACHE_IMAGE": self.model.config["cache-image"],
}
return config_json

def _check_container_connection(self, container: Container) -> None:
"""Check if connection can be made with container.
Args:
container: the named container in a unit to check.

Raises:
ErrorWithStatus if the connection cannot be made.
"""
if not container.can_connect():
raise ErrorWithStatus("Pod startup is not complete", MaintenanceStatus)
return env_vars

def _check_status(self):
"""Check status of workload and set status accordingly."""
Expand All @@ -301,28 +276,6 @@ def _check_status(self):
raise ErrorWithStatus("Workload failed health check", MaintenanceStatus)
self.model.unit.status = ActiveStatus()

def _upload_files_to_container(self, config_json):
"""Upload required files to container."""
try:
self._check_container_connection(self.container)
except ErrorWithStatus as error:
self.model.unit.status = error.status
raise error
try:
with open("src/sample_config.json", "r") as sample_config:
file_content = sample_config.read()
self.container.push(SAMPLE_CONFIG, file_content, make_dirs=True)
except ErrorWithStatus as error:
self.logger.error("Failed to upload sample config to container.")
raise error
try:
file_content = json.dumps(config_json)
config = CONFIG_DIR / "config.json"
self.container.push(config, file_content, make_dirs=True)
except ErrorWithStatus as error:
self.logger.error("Failed to upload config to container.")
raise error

def _send_info(self, interfaces):
if interfaces["kfp-api"]:
interfaces["kfp-api"].send_data(
Expand Down Expand Up @@ -680,12 +633,9 @@ def _on_event(self, event, force_conflicts: bool = False) -> None:
# Set up all relations/fetch required data
try:
self._check_leader()
interfaces = self._get_interfaces()
config_json = self._generate_config(interfaces)
self._upload_files_to_container(config_json)
self._apply_k8s_resources(force_conflicts=force_conflicts)
update_layer(self._container_name, self._container, self._kfp_api_layer, self.logger)
self._send_info(interfaces)
self._send_info(self._get_interfaces())
except ErrorWithStatus as err:
self.model.unit.status = err.status
self.logger.error(f"Failed to handle {event} with error: {err}")
Expand Down
12 changes: 0 additions & 12 deletions charms/kfp-api/src/sample_config.json
View file

This file was deleted.

87 changes: 60 additions & 27 deletions charms/kfp-api/tests/unit/test_operator.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,7 @@ def test_mysql_relation(
expected_returned_data,
expected_raises,
expected_status,
mocked_lightkube_client,
harness: Harness,
):
harness.set_leader(True)
Expand All @@ -132,7 +133,7 @@ def test_mysql_relation(
harness.charm._get_db_data()

@patch("charm.KubernetesServicePatch", lambda x, y: None)
def test_mysql_relation_too_many_relations(self, harness: Harness):
def test_mysql_relation_too_many_relations(self, mocked_lightkube_client, harness: Harness):
harness.set_leader(True)
harness.begin()
harness.container_pebble_ready(KFP_API_CONTAINER_NAME)
Expand All @@ -150,7 +151,7 @@ def test_mysql_relation_too_many_relations(self, harness: Harness):
)

@patch("charm.KubernetesServicePatch", lambda x, y: None)
def test_kfp_viz_relation_missing(self, harness: Harness):
def test_kfp_viz_relation_missing(self, mocked_lightkube_client, harness: Harness):
harness.set_leader(True)
harness.begin()
harness.container_pebble_ready(KFP_API_CONTAINER_NAME)
Expand Down Expand Up @@ -309,6 +310,7 @@ def test_relations_that_provide_data(
expected_returned_data,
expected_raises,
expected_status,
mocked_lightkube_client,
harness: Harness,
):
harness.set_leader(True)
Expand Down Expand Up @@ -360,31 +362,33 @@ def test_install_with_all_inputs_and_pebble(
harness.update_relation_data(mysql_rel_id, "mysql-provider/0", mysql_data)

# object storage relation
os_data = {
objectstorage_data = {
"access-key": "access-key",
"namespace": "namespace",
"port": 1234,
"secret-key": "secret-key",
"secure": True,
"service": "service",
}
objectstorage_data_dict = {
"_supported_versions": "- v1",
"data": yaml.dump(
{
"access-key": "access-key",
"namespace": "namespace",
"port": 1234,
"secret-key": "secret-key",
"secure": True,
"service": "service",
}
),
"data": yaml.dump(objectstorage_data),
}
os_rel_id = harness.add_relation("object-storage", "storage-provider")
harness.add_relation_unit(os_rel_id, "storage-provider/0")
harness.update_relation_data(os_rel_id, "storage-provider", os_data)
objectstorage_rel_id = harness.add_relation("object-storage", "storage-provider")
harness.add_relation_unit(objectstorage_rel_id, "storage-provider/0")
harness.update_relation_data(
objectstorage_rel_id, "storage-provider", objectstorage_data_dict
)

# kfp-viz relation
kfp_viz_data = {
"_supported_versions": "- v1",
"data": yaml.dump({"service-name": "unset", "service-port": "1234"}),
"service-name": "viz-service",
"service-port": "1234",
}
kfp_viz_data_dict = {"_supported_versions": "- v1", "data": yaml.dump(kfp_viz_data)}
kfp_viz_id = harness.add_relation("kfp-viz", "kfp-viz")
harness.add_relation_unit(kfp_viz_id, "kfp-viz/0")
harness.update_relation_data(kfp_viz_id, "kfp-viz", kfp_viz_data)
harness.update_relation_data(kfp_viz_id, "kfp-viz", kfp_viz_data_dict)

# example kfp-api provider relation
kfpapi_data = {
Expand Down Expand Up @@ -431,22 +435,53 @@ def test_install_with_all_inputs_and_pebble(
"-logtostderr=true "
)
assert pebble_exec_command == f"bash -c '{exec_command}'"

expected_env = {
"AUTO_UPDATE_PIPELINE_DEFAULT_VERSION": harness.charm.config[
"auto-update-default-version"
],
"KFP_API_SERVICE_NAME": KFP_API_SERVICE_NAME,
"KUBEFLOW_USERID_HEADER": "kubeflow-userid",
"KUBEFLOW_USERID_PREFIX": "",
"POD_NAMESPACE": harness.charm.model.name,
"OBJECTSTORECONFIG_SECURE": "false",
"OBJECTSTORECONFIG_BUCKETNAME": harness.charm.config["object-store-bucket-name"],
"DBCONFIG_USER": "root",
"DBCONFIG_PASSWORD": mysql_data["root_password"],
"DBCONFIG_DBNAME": mysql_data["database"],
"DBCONFIG_HOST": mysql_data["host"],
"DBCONFIG_PORT": mysql_data["port"],
"DBCONFIG_CONMAXLIFETIME": "120s",
"DB_DRIVER_NAME": "mysql",
"DBCONFIG_MYSQLCONFIG_USER": "root",
"DBCONFIG_MYSQLCONFIG_PASSWORD": mysql_data["root_password"],
"DBCONFIG_MYSQLCONFIG_DBNAME": mysql_data["database"],
"DBCONFIG_MYSQLCONFIG_HOST": mysql_data["host"],
"DBCONFIG_MYSQLCONFIG_PORT": mysql_data["port"],
"OBJECTSTORECONFIG_ACCESSKEY": objectstorage_data["access-key"],
"OBJECTSTORECONFIG_SECRETACCESSKEY": objectstorage_data["secret-key"],
"DEFAULTPIPELINERUNNERSERVICEACCOUNT": "default-editor",
"MULTIUSER": "true",
"VISUALIZATIONSERVICE_NAME": kfp_viz_data["service-name"],
"VISUALIZATIONSERVICE_PORT": kfp_viz_data["service-port"],
"ML_PIPELINE_VISUALIZATIONSERVER_SERVICE_HOST": kfp_viz_data["service-name"],
"ML_PIPELINE_VISUALIZATIONSERVER_SERVICE_PORT": kfp_viz_data["service-port"],
"CACHE_IMAGE": harness.charm.config["cache-image"],
}
test_env = pebble_plan_info["services"][KFP_API_SERVICE_NAME]["environment"]
# there should be 1 environment variable
assert 1 == len(test_env)

assert test_env == expected_env
assert "test_model" == test_env["POD_NAMESPACE"]

@patch("charm.KubernetesServicePatch", lambda x, y: None)
@patch("charm.KfpApiOperator._apply_k8s_resources")
@patch("charm.KfpApiOperator._check_status")
@patch("charm.KfpApiOperator._generate_config")
@patch("charm.KfpApiOperator._upload_files_to_container")
@patch("charm.KfpApiOperator._generate_environment")
def test_update_status(
self,
_apply_k8s_resources: MagicMock,
_check_status: MagicMock,
_generate_config: MagicMock,
_upload_files_to_conainer: MagicMock,
_generate_environment: MagicMock,
harness: Harness,
):
"""Test update status handler."""
Expand All @@ -456,11 +491,9 @@ def test_update_status(

# test successful update status
_apply_k8s_resources.reset_mock()
_upload_files_to_conainer.reset_mock()
harness.charm.on.update_status.emit()
# this will enforce the design in which main event handler is executed in update-status
_apply_k8s_resources.assert_called()
_upload_files_to_conainer.assert_called()
# check status should be called
_check_status.assert_called()

Expand Down

0 comments on commit 468b6e7

Please sign in to comment.