feat(launch): support volume mounts and security contexts in kubernetes runner

tests/unit_tests_old/tests_launch/test_launch_kubernetes.py

@@ -247,6 +247,7 @@ def test_launch_kube(
                 "node_name": "test-node-name",
                 "node_selectors": {"test-selector": "test-value"},
                 "tolerations": [{"key": "test-key", "value": "test-value"}],
+                "volumes": [{"name": "test-volume", "host_path": "/test/path"}],
             },
         },
     }
@@ -272,6 +273,7 @@ def test_launch_kube(
     assert job.spec.template.spec.preemption_policy == args["preemption_policy"]
     assert job.spec.template.spec.node_name == args["node_name"]
     assert job.spec.template.spec.tolerations == args["tolerations"]
+    assert job.spec.template.spec.volumes == args["volumes"]
     assert (
         job.spec.template.spec.node_selector["test-selector"]
         == args["node_selectors"]["test-selector"]

wandb/sdk/launch/runner/kubernetes.py

@@ -169,6 +169,8 @@ def populate_pod_spec(
             pod_spec["nodeSelectors"] = resource_args.get("node_selectors")
         if resource_args.get("tolerations"):
             pod_spec["tolerations"] = resource_args.get("tolerations")
+        if resource_args.get("volumes"):
+            pod_spec["volumes"] = resource_args.get("volumes")
 
     def populate_container_resources(
         self, containers: List[Dict[str, Any]], resource_args: Dict[str, Any]
@@ -365,6 +367,7 @@ def run(
         for cont in containers:
             cont["env"] = [{"name": k, "value": v} for k, v in merged_env_vars.items()]
         pod_spec["containers"] = containers
+
         pod_template["spec"] = pod_spec
         pod_template["metadata"] = pod_metadata
         if secret is not None: