Solve an issue with MP container checkouts

An issue was identified where the _cont_inst property was being stripped from
the modified Host class during MP checkouts.
The solution I opted for was to recreate the missing container instance
post-pickle. This removed the entire need for the _cont_inst property.
I then took this further with additional pickle protection.

broker/helpers.py

@@ -5,6 +5,7 @@
 import json
 import logging
 import os
+import pickle
 import sys
 import tarfile
 import time
@@ -548,3 +549,46 @@ def temporary_tar(paths):
             tar.add(path, arcname=path.name)
     yield temp_tar.absolute()
     temp_tar.unlink()
+
+
+class PickleSafe:
+    """A class that helps with pickling and unpickling complex objects"""
+
+    def _pre_pickle(self):
+        """This method is called before pickling an object"""
+        pass
+
+    def _post_pickle(self, purified):
+        """This method is called after pickling an object
+
+        purified will be a list of names of attributes that were removed
+        """
+        pass
+
+    def _purify(self):
+        """Strip all unpickleable attributes from a Host before pickling"""
+        self.purified = getattr(self, "purified", [])
+        for name in list(self.__dict__):
+            self._purify_target = name
+            try:
+                pickle.dumps(self.__dict__[name])
+            except (pickle.PicklingError, AttributeError):
+                self.__dict__[name] = None
+                self.purified.append(name)
+
+    def __getstate__(self):
+        """If a session is active, remove it for pickle compatability"""
+        self._pre_pickle()
+        try:
+            self._purify()
+        except RecursionError:
+                logger.warning(f"Recursion limit reached on {self._purify_target}")
+                self.__dict__[self._purify_target] = None
+                self.__getstate__()
+        self.__dict__.pop("_purify_target", None)
+        return self.__dict__
+
+    def __setstate__(self, state):
+        """Sometimes pickle strips things that we need. This should be used to restore them"""
+        self.__dict__.update(state)
+        self._post_pickle(purified=getattr(self, "purified", []))

broker/hosts.py

@@ -1,12 +1,12 @@
 # from functools import cached_property
-import pickle
 from logzero import logger
 from broker.exceptions import NotImplementedError
+from broker.helpers import PickleSafe
 from broker.session import ContainerSession, Session
 from broker.settings import settings
 
 
-class Host:
+class Host(PickleSafe):
 
     default_timeout = 0  # timeout in ms, 0 is infinite
 
@@ -43,26 +43,12 @@ def session(self):
                 self.connect()
         return self._session
 
-    def __getstate__(self):
-        """If a session is active, remove it for pickle compatability"""
+    def _pre_pickle(self):
         self.close()
-        try:
-            self._purify()
-        except RecursionError:
-                logger.warning(f"Recursion limit reached on {self._purify_target}")
-                self.__dict__[self._purify_target] = None
-                self.__getstate__()
-        self.__dict__.pop("_purify_target", None)
-        return self.__dict__
-
-    def _purify(self):
-        """Strip all unpickleable attributes from a Host before pickling"""
-        for name in list(self.__dict__):
-            self._purify_target = name
-            try:
-                pickle.dumps(self.__dict__[name])
-            except (pickle.PicklingError, AttributeError):
-                self.__dict__[name] = None
+
+    def _post_pickle(self, purified):
+        if "_cont_inst" in purified and not getattr(self, "_checked_in", False):
+            self._cont_inst = self._prov_inst._cont_inst_by_name(self.name)
 
     def connect(
         self, username=None, password=None, timeout=None, port=22, key_filename=None

broker/providers/__init__.py

@@ -2,11 +2,12 @@
 import dynaconf
 
 from broker import exceptions
+from broker.helpers import PickleSafe
 from broker.settings import settings
 from logzero import logger
 
 
-class Provider:
+class Provider(PickleSafe):
     # Populate with a list of Dynaconf Validators specific to your provider
     _validators = []
     # Set to true if you don't want your provider shown in the CLI
@@ -33,6 +34,7 @@ def _validate_settings(self, instance_name=None):
 
         :param instance_name: A string matching an instance name
         """
+        instance_name = instance_name or getattr(self, "instance", None)
         section_name = self.__class__.__name__.upper()
         # make sure each instance isn't loading values from another
         fresh_settings = settings.get_fresh(section_name)

broker/providers/container.py

@@ -22,21 +22,14 @@ def container_info(container_inst):
     }
 
 
-@property
-def _cont_inst(self):
-    """Returns a live container object instance"""
-    if not getattr(self, "_cont_inst_p", None):
-        self._cont_inst_p = self._prov_inst._cont_inst_by_name(self.name)
-    return self._cont_inst_p
-
-
 def _host_release():
     caller_host = inspect.stack()[1][0].f_locals["host"]
-    if not caller_host._cont_inst_p:
-        caller_host._cont_inst_p = caller_host._prov_inst._cont_inst_by_name(
+    if not caller_host._cont_inst:
+        caller_host._cont_inst = caller_host._prov_inst._cont_inst_by_name(
             caller_host.name
         )
-    caller_host._cont_inst_p.remove(v=True, force=True)
+    caller_host._cont_inst.remove(v=True, force=True)
+    caller_host._checked_in = True
 
 
 class Container(Provider):
@@ -80,21 +73,25 @@ def __init__(self, **kwargs):
                 "Container",
                 f"Broker has no bind for {settings.container.runtime} containers",
             )
-        self._runtime = None  # this will be used later
+        self.runtime = self._runtime_cls(
+            host=settings.container.host,
+            username=settings.container.host_username,
+            password=settings.container.host_password,
+            port=settings.container.host_port,
+            timeout=settings.container.timeout,
+        )
         self._name_prefix = settings.container.get("name_prefix", getpass.getuser())
 
-    @property
-    def runtime(self):
-        """Making this a property helps to recover from pickle environments"""
-        if not self._runtime:
-            self._runtime = self._runtime_cls(
-                host=settings.container.host,
-                username=settings.container.host_username,
-                password=settings.container.host_password,
-                port=settings.container.host_port,
-                timeout=settings.container.timeout,
-            )
-        return self._runtime
+
+    def _post_pickle(self, purified):
+        self._validate_settings()
+        self.runtime = self._runtime_cls(
+            host=settings.container.host,
+            username=settings.container.host_username,
+            password=settings.container.host_password,
+            port=settings.container.host_port,
+            timeout=settings.container.timeout,
+        )
 
     def _ensure_image(self, name):
         """Check if an image exists on the provider, attempt a pull if not"""
@@ -128,14 +125,13 @@ def _set_attributes(self, host_inst, broker_args=None, cont_inst=None):
         host_inst.__dict__.update(
             {
                 "_prov_inst": self,
-                "_cont_inst_p": cont_inst,
+                "_cont_inst": cont_inst,
                 "_broker_provider": "Container",
                 "_broker_provider_instance": self.instance,
                 "_broker_args": broker_args,
                 "release": _host_release,
             }
         )
-        host_inst.__class__._cont_inst = _cont_inst
 
     def _port_mapping(self, image, **kwargs):
         """

tests/data/cli_scenarios/satlab/checkout_rhel78.yaml

@@ -1,2 +1,2 @@
 workflow: deploy-base-rhel
-rhel_version: "7.8"
+deploy_rhel_version: "7.8"

tests/functional/test_containers.py

@@ -75,3 +75,15 @@ def test_container_e2e():
         c_host.session.sftp_write("broker_settings.yaml", "/root")
         res = c_host.execute("ls")
         assert "broker_settings.yaml" in res.stdout
+
+
+def test_container_e2e_mp():
+    with Broker(container_host="ubi8:latest", _count=2) as c_hosts:
+        for c_host in c_hosts:
+            assert c_host._cont_inst.top()['Processes']
+            res = c_host.execute("hostname")
+            assert res.stdout.strip() == c_host.hostname
+            # Test that a file can be uploaded to the container
+            c_host.session.sftp_write("broker_settings.yaml", "/root")
+            res = c_host.execute("ls")
+            assert "broker_settings.yaml" in res.stdout