diff --git a/calrissian/context.py b/calrissian/context.py index 9b023f0..4dd6124 100644 --- a/calrissian/context.py +++ b/calrissian/context.py @@ -23,4 +23,6 @@ def __init__(self, kwargs=None): self.pod_serviceaccount = None self.tool_logs_basepath = None self.max_gpus = None + self.no_network_access_pod_labels = None + self.network_access_pod_labels = None return super(CalrissianRuntimeContext, self).__init__(kwargs) diff --git a/calrissian/job.py b/calrissian/job.py index d6cc300..dd7b20b 100644 --- a/calrissian/job.py +++ b/calrissian/job.py @@ -198,8 +198,10 @@ def add_emptydir_volume_binding(self, name, target): class KubernetesPodBuilder(object): - def __init__(self, name, container_image, environment, volume_mounts, volumes, command_line, stdout, stderr, stdin, resources, labels, nodeselectors, security_context, serviceaccount, requirements=None, hints=None): + def __init__(self, name, builder, container_image, environment, volume_mounts, volumes, command_line, stdout, stderr, stdin, labels, nodeselectors, security_context, serviceaccount, no_network_access_pod_labels=None, network_access_pod_labels=None): self.name = name + self.builder = builder + self.cwl_version = self.builder.cwlVersion self.container_image = container_image self.environment = environment self.volume_mounts = volume_mounts @@ -208,13 +210,15 @@ def __init__(self, name, container_image, environment, volume_mounts, volumes, c self.stdout = stdout self.stderr = stderr self.stdin = stdin - self.resources = resources + self.resources = self.builder.resources self.labels = labels self.nodeselectors = nodeselectors self.security_context = security_context self.serviceaccount = serviceaccount - self.requirements = {} if requirements is None else requirements - self.hints = [] if hints is None else hints + self.no_network_access_pod_labels = no_network_access_pod_labels + self.network_access_pod_labels = network_access_pod_labels + self.requirements = {} if self.builder.requirements is None else self.builder.requirements + self.hints = [] if self.builder.hints is None else self.builder.hints def pod_name(self): tag = random_tag() @@ -340,6 +344,21 @@ def pod_labels(self): Submitted labels must be strings :return: """ + if self.cwl_version in ["v1.0"]: + network_access = True + else: + network_access = False + + for requirement in self.requirements: + if "class" in requirement.keys() and requirement["class"] in ["NetworkAccess"]: + network_access = True if requirement.get("networkAccess") == "true" else False + break + if not network_access and self.no_network_access_pod_labels: + self.labels = {**self.labels, **self.no_network_access_pod_labels} + + if network_access and self.network_access_pod_labels: + self.labels = {**self.labels, **self.network_access_pod_labels} + return {str(k): str(v) for k, v in self.labels.items()} def pod_nodeselectors(self): @@ -515,7 +534,19 @@ def get_pod_labels(self, runtimeContext): return read_yaml(runtimeContext.pod_labels) else: return {} - + + def get_network_access_pod_labels(self, runtimeContext): + if runtimeContext.network_access_pod_labels: + return read_yaml(runtimeContext.network_access_pod_labels) + else: + return {} + + def get_no_network_access_pod_labels(self, runtimeContext): + if runtimeContext.no_network_access_pod_labels: + return read_yaml(runtimeContext.no_network_access_pod_labels) + else: + return {} + def get_pod_nodeselectors(self, runtimeContext): if runtimeContext.pod_nodeselectors: return read_yaml(runtimeContext.pod_nodeselectors) @@ -576,6 +607,7 @@ def create_kubernetes_runtime(self, runtimeContext): k8s_builder = KubernetesPodBuilder( self.name, + self.builder, self._get_container_image(), self.environment, self.volume_builder.volume_mounts, @@ -584,13 +616,12 @@ def create_kubernetes_runtime(self, runtimeContext): self.stdout, self.stderr, self.stdin, - self.builder.resources, self.get_pod_labels(runtimeContext), self.get_pod_nodeselectors(runtimeContext), self.get_security_context(runtimeContext), self.get_pod_serviceaccount(runtimeContext), - self.builder.requirements, - self.builder.hints + self.get_no_network_access_pod_labels(runtimeContext), + self.get_network_access_pod_labels(runtimeContext), ) built = k8s_builder.build() log.debug('{}\n{}{}\n'.format('-' * 80, yaml.dump(built), '-' * 80)) diff --git a/calrissian/main.py b/calrissian/main.py index 6f27aa4..8fa4e2f 100644 --- a/calrissian/main.py +++ b/calrissian/main.py @@ -49,6 +49,8 @@ def add_arguments(parser): parser.add_argument('--stderr', type=Text, nargs='?', help='Output file name to tee standard error to (includes tool logs)') parser.add_argument('--tool-logs-basepath', type=Text, nargs='?', help='Base path for saving the tool logs') parser.add_argument('--conf', help='Defines the default values for the CLI arguments', action='append') + parser.add_argument('--no-network-access-pod-label', type=Text, nargs='?', help='YAML file to set the pod label to use for disabling network access') + parser.add_argument('--network-access-pod-label', type=Text, nargs='?', help='YAML file to set the pod label to use for enabling network access') def print_version(): diff --git a/pyproject.toml b/pyproject.toml index e9b5997..9fea01e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ name = "calrissian" dynamic = ["version"] description = 'CWL runner for Kubernetes' readme = "README.md" -requires-python = ">=3.8" +requires-python = ">=3.9" license = "MIT" keywords = [] authors = [ @@ -21,24 +21,26 @@ classifiers = [ "Environment :: Console", "Intended Audience :: Developers", "Programming Language :: Python", - "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", "Programming Language :: Python :: Implementation :: CPython", "Programming Language :: Python :: Implementation :: PyPy", ] dependencies = [ "urllib3==1.26.18", "kubernetes==28.1.0", - "cwltool==3.1.20240708091337", + "cwltool==3.1.20250110105449", "tenacity==8.2.3", "importlib-metadata==6.8.0", - "msgpack==1.0.7", + "msgpack>=1.0.7", "typing-extensions==4.8.0", - "freezegun==1.2.2", - "setuptools==70.0.0" + "freezegun==1.5.1", + "setuptools==70.0.0", + "shellescape==3.8.1", + "python-dateutil>=2.7" ] [project.urls] @@ -81,13 +83,15 @@ skip-install = false dependencies = [ "urllib3==1.26.18", "kubernetes==28.1.0", - "cwltool==3.1.20240708091337", + "cwltool==3.1.20250110105449", "tenacity==8.2.3", "importlib-metadata==6.8.0", - "msgpack==1.0.7", + "msgpack>=1.0.7", "typing-extensions==4.8.0", - "freezegun==1.2.2", - "setuptools==70.0.0" + "freezegun==1.5.1", + "setuptools==70.0.0", + "shellescape==3.8.1", + "python-dateutil>=2.7" ] [tool.hatch.envs.prod] @@ -106,13 +110,15 @@ dependencies = [ "coverage", "urllib3==1.26.18", "kubernetes==28.1.0", - "cwltool==3.1.20240708091337", + "cwltool==3.1.20250110105449", "tenacity==8.2.3", "importlib-metadata==6.8.0", - "msgpack==1.0.7", + "msgpack>=1.0.7", "typing-extensions==4.8.0", - "freezegun==1.2.2", - "setuptools==70.0.0" + "freezegun==1.5.1", + "setuptools==70.0.0", + "shellescape==3.8.1", + "python-dateutil>=2.7" ] [tool.hatch.envs.test.env-vars] @@ -124,7 +130,7 @@ testv = "hatch run nose2 --verbose" cov = ["coverage run --source=calrissian -m nose2", "coverage report"] [[tool.hatch.envs.test.matrix]] -python = ["3.8", "3.9", "3.10", "3.11", "3.12"] +python = ["3.9", "3.10", "3.11", "3.12", "3.13"] [tool.hatch.envs.docs] skip-install = false diff --git a/tests/test_job.py b/tests/test_job.py index 50f6bf9..0551cde 100644 --- a/tests/test_job.py +++ b/tests/test_job.py @@ -273,7 +273,12 @@ def test_add_persistent_volume_entries_from_pod_preserves_readonly(self, mock_ku class KubernetesPodBuilderTestCase(TestCase): def setUp(self): + builder = Mock() + builder.cwlVersion = "v1.2" + builder.requirements = [] + builder.resources = {'cores': 1, 'ram': 1024} self.name = 'PodName' + self.builder = builder self.container_image = 'dockerimage:1.0' self.environment = {'K1':'V1', 'K2':'V2', 'HOME': '/homedir'} self.volume_mounts = [Mock(), Mock()] @@ -282,14 +287,16 @@ def setUp(self): self.stdout = 'stdout.txt' self.stderr = 'stderr.txt' self.stdin = 'stdin.txt' - self.resources = {'cores': 1, 'ram': 1024} self.labels = {'key1': 'val1', 'key2': 123} self.nodeselectors = {'disktype': 'ssd', 'cachelevel': 2} self.security_context = { 'runAsUser': os.getuid(),'runAsGroup': os.getgid() } self.pod_serviceaccount = "podmanager" - self.pod_builder = KubernetesPodBuilder(self.name, self.container_image, self.environment, self.volume_mounts, + self.no_network_access_pod_labels = {"calrissian-network": "disabled"} + self.network_access_pod_labels = {"calrissian-network": "enabled"} + self.pod_builder = KubernetesPodBuilder(self.name, self.builder, self.container_image, self.environment, self.volume_mounts, self.volumes, self.command_line, self.stdout, self.stderr, self.stdin, - self.resources, self.labels, self.nodeselectors, self.security_context, self.pod_serviceaccount) + self.labels, self.nodeselectors, self.security_context, self.pod_serviceaccount, + self.no_network_access_pod_labels, self.network_access_pod_labels) @patch('calrissian.job.random_tag') def test_safe_pod_name(self, mock_random_tag): @@ -375,14 +382,37 @@ def test_gpu_hints(self): } self.assertEqual(expected, resources) + def test_network_access_1_2(self): + self.pod_builder.cwl_version = "v1.2" + self.pod_builder.requirements = [OrderedDict([("class", "NetworkAccess"), ("networkAccess", "true")])] + self.assertEqual(self.pod_builder.pod_labels(), {"calrissian-network": "enabled", 'key1':'val1', 'key2':'123'}) + + def test_no_network_access_1_2(self): + self.pod_builder.cwl_version = "v1.2" + self.pod_builder.requirements = [OrderedDict([("class", "NetworkAccess"), ("networkAccess", "false")])] + self.assertEqual(self.pod_builder.pod_labels(), {"calrissian-network": "disabled", 'key1':'val1', 'key2':'123'}) + + def test_network_access_1_0(self): + self.pod_builder.cwl_version = "v1.0" + self.pod_builder.requirements = [OrderedDict([])] + self.assertEqual(self.pod_builder.pod_labels(), {"calrissian-network": "enabled", 'key1':'val1', 'key2':'123'}) + def test_string_labels(self): self.pod_builder.labels = {'key1': 123} - self.assertEqual(self.pod_builder.pod_labels(), {'key1':'123'}) + self.assertEqual(self.pod_builder.pod_labels(), {"calrissian-network": "disabled", 'key1':'123'}) def test_string_nodeselectors(self): self.pod_builder.nodeselectors = {'cachelevel': 2} self.assertEqual(self.pod_builder.pod_nodeselectors(), {'cachelevel':'2'}) + def test_string_no_network_access_pod_label(self): + self.pod_builder.no_network_access_pod_labels = {"calrissian-network": "disabled"} + self.assertEqual(self.pod_builder.pod_labels(), {"calrissian-network": "disabled", 'key1': 'val1', 'key2': '123'}) + + def test_string_network_access_pod_label(self): + self.pod_builder.network_access_pod_labels = {"calrissian-network": "enabled"} + self.assertEqual(self.pod_builder.pod_labels(), {"calrissian-network": "disabled", 'key1': 'val1', 'key2': '123'}) + def test_init_containers_empty_when_no_stdout_or_stderr(self): self.pod_builder.stdout = None self.pod_builder.stderr = None @@ -428,6 +458,7 @@ def test_build(self, mock_random_tag): 'labels': { 'key1': 'val1', 'key2': '123', + 'calrissian-network': 'disabled', } }, 'apiVersion': 'v1', @@ -657,6 +688,7 @@ def realpath(path): # creates a KubernetesPodBuilder self.assertEqual(mock_pod_builder.call_args, call( job.name, + job.builder, job._get_container_image(), job.environment, job.volume_builder.volume_mounts, @@ -665,13 +697,12 @@ def realpath(path): job.stdout, job.stderr, job.stdin, - job.builder.resources, mock_read_yaml.return_value, mock_read_yaml.return_value, job.get_security_context(mock_runtime_context), None, - job.builder.requirements, - job.builder.hints, + mock_read_yaml.return_value, + mock_read_yaml.return_value, )) # calls builder.build # returns that diff --git a/tests/test_main.py b/tests/test_main.py index e925f8d..875fc54 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -63,7 +63,7 @@ def test_main_calls_cwlmain_returns_exit_code(self, mock_activate_logging, mock_ def test_add_arguments(self): mock_parser = Mock() add_arguments(mock_parser) - self.assertEqual(mock_parser.add_argument.call_count, 12) + self.assertEqual(mock_parser.add_argument.call_count, 14) @patch('calrissian.main.sys') def test_parse_arguments_exits_without_ram_or_cores(self, mock_sys):