Merge remote-tracking branch 'origin/master' into databricks-fix

granulate_utils/containers/cri.py

@@ -17,8 +17,9 @@
 
 import json
 from contextlib import contextmanager, suppress
+from dataclasses import dataclass, field
 from datetime import datetime, timezone
-from typing import Any, List, Optional, Type, TypeVar
+from typing import Any, Dict, List, Optional, Type, TypeVar
 
 import grpc  # type: ignore # no types-grpc sadly
 import psutil
@@ -122,14 +123,15 @@ def _create_container(
             with suppress(psutil.NoSuchProcess):
                 process = psutil.Process(pid)
 
-        return Container(
+        return K8sContainer(
             runtime=self.runtime_name,
             name=self._reconstruct_name(container),
             id=container.id,
             labels=container.labels,
             running=container.state == self.api.api_pb2.CONTAINER_RUNNING,
             process=process,
             time_info=time_info,
+            annotations=container.annotations,
         )
 
 
@@ -182,3 +184,8 @@ def get_container(self, container_id: str, all_info: bool) -> Container:
 
     def get_runtimes(self) -> List[str]:
         return [client.runtime_name for client in self._clients]
+
+
+@dataclass
+class K8sContainer(Container):
+    annotations: Dict[str, str] = field(default_factory=dict)

granulate_utils/linux/ns.py

@@ -246,27 +246,13 @@ def run_in_ns(
 
     for ns in nstypes:
         assert_ns_str(ns)
-    # make sure "mnt" is last, once we change it our /proc is gone
-    nstypes = sorted(nstypes, key=lambda ns: 1 if ns == "mnt" else 0)
 
     ret: Union[T, _Sentinel] = _SENTINEL
     exc: Optional[BaseException] = None
 
     def _switch_and_run():
         try:
-            global libc
-            if libc is None:
-                libc = ctypes.CDLL("libc.so.6")
-
-            for nstype in nstypes:
-                if not is_same_ns(target_pid, nstype):
-                    flag = NsType[nstype].value
-                    if libc.unshare(flag) != 0:
-                        raise ValueError(f"Failed to unshare({nstype})")
-
-                    with open(f"/proc/{target_pid}/ns/{nstype}", "r") as nsf:
-                        if libc.setns(nsf.fileno(), flag) != 0:
-                            raise ValueError(f"Failed to setns({nstype}) (to pid {target_pid})")
+            enter_process_ns(nstypes, target_pid)
 
             nonlocal ret
             ret = callback()
@@ -363,3 +349,30 @@ def get_host_pid(nspid: int, container_id: str) -> Optional[int]:
             continue
 
     return None
+
+
+def enter_process_ns(
+    nstypes: List[str],
+    target_pid: int = 1,
+) -> None:
+    """
+    Swaps to a set of the namespaces of a target process.
+    """
+    for ns in nstypes:
+        assert_ns_str(ns)
+    # make sure "mnt" is last, once we change it our /proc is gone
+    nstypes = sorted(nstypes, key=lambda ns: 1 if ns == "mnt" else 0)
+
+    global libc
+    if libc is None:
+        libc = ctypes.CDLL("libc.so.6")
+
+    for nstype in nstypes:
+        if not is_same_ns(target_pid, nstype):
+            flag = NsType[nstype].value
+            if libc.unshare(flag) != 0:
+                raise ValueError(f"Failed to unshare({nstype})")
+
+            with open(f"/proc/{target_pid}/ns/{nstype}", "r") as nsf:
+                if libc.setns(nsf.fileno(), flag) != 0:
+                    raise ValueError(f"Failed to setns({nstype}) (to pid {target_pid})")