Python pickle.__name__() Examples

def test_load_pyfunc_loads_torch_model_using_pickle_module_specified_at_save_time(
        module_scoped_subclassed_model, model_path):
    custom_pickle_module = pickle

    mlflow.pytorch.save_model(
        path=model_path,
        pytorch_model=module_scoped_subclassed_model,
        conda_env=None,
        pickle_module=custom_pickle_module)

    import_module_fn = importlib.import_module
    imported_modules = []

    def track_module_imports(module_name):
        imported_modules.append(module_name)
        return import_module_fn(module_name)

    with mock.patch("importlib.import_module") as import_mock,\
            mock.patch("torch.load") as torch_load_mock:
        import_mock.side_effect = track_module_imports
        pyfunc.load_pyfunc(model_path)

    torch_load_mock.assert_called_with(mock.ANY, pickle_module=custom_pickle_module)
    assert custom_pickle_module.__name__ in imported_modules 

def get(cls, curve):
            if isinstance(curve, CurveDB):
                return curve
            elif isinstance(curve, list):
                return [CurveDB.get(c) for c in curve]
            else:
                with open(os.path.join(CurveDB._dirname, str(curve) + cls.file_extension),
                          'rb' if file_backend.__name__ == 'pickle' else 'r')\
                        as f:
                    # rb is for compatibility with python 3
                    # see http://stackoverflow.com/questions/5512811/builtins-typeerror-must-be-str-not-bytes
                    curve = CurveDB()
                    curve._pk, curve.params, data = file_backend.load(f)
                    curve.data = tuple([np.asarray(a) for a in data])
                if isinstance(curve.data, pd.Series):  # for backwards compatibility
                    x, y = curve.data.index.values, curve.data.values
                    curve.data = (x, y)
                return curve 

def __repr__(self):
        if hasattr(self, '__getstate__'):
            state = repr(self.__getstate__())
        else:
            state = sorted_dict_repr(self.__dict__)
        return "<%s instance with state %s>" % (
            type(self).__name__, normalized_repr(state)) 

def test_load_negative(self):
        if cPickle.__name__ == "cPickle":   # pickle vs. cPickle report different exceptions, even on Cpy
            filename = os.tempnam()
            for temp in ['\x02', "No"]:
                self.write_to_file(filename, content=temp)
                f = open(filename)
                self.assertRaises(cPickle.UnpicklingError, cPickle.load, f)
                f.close() 

def __repr__(self):
        if hasattr(self, '__getstate__'):
            state = repr(self.__getstate__())
        else:
            state = sorted_dict_repr(self.__dict__)
        return "<%s instance with state %s>" % (
            type(self).__name__, normalized_repr(state)) 

def test_load_negative(self):
        if cPickle.__name__ == "_pickle":   # pickle vs. cPickle report different exceptions, even on Cpy
            filename = os.tempnam()
            for temp in ['\x02', "No"]:
                self.write_to_file(filename, content=temp)
                f = open(filename)
                self.assertRaises(cPickle.UnpicklingError, cPickle.load, f)
                f.close() 

def test_load_model_loads_torch_model_using_pickle_module_specified_at_save_time(
        module_scoped_subclassed_model):
    custom_pickle_module = pickle

    artifact_path = "pytorch_model"
    with mlflow.start_run():
        mlflow.pytorch.log_model(
            artifact_path=artifact_path,
            pytorch_model=module_scoped_subclassed_model,
            conda_env=None,
            pickle_module=custom_pickle_module)
        model_uri = "runs:/{run_id}/{artifact_path}".format(
            run_id=mlflow.active_run().info.run_id,
            artifact_path=artifact_path)

    import_module_fn = importlib.import_module
    imported_modules = []

    def track_module_imports(module_name):
        imported_modules.append(module_name)
        return import_module_fn(module_name)

    with mock.patch("importlib.import_module") as import_mock,\
            mock.patch("torch.load") as torch_load_mock:
        import_mock.side_effect = track_module_imports
        pyfunc.load_pyfunc(model_uri=model_uri)

    torch_load_mock.assert_called_with(mock.ANY, pickle_module=custom_pickle_module)
    assert custom_pickle_module.__name__ in imported_modules 

def test_load_model_allows_user_to_override_pickle_module_via_keyword_argument(
        module_scoped_subclassed_model, model_path):
    mlflow.pytorch.save_model(
        path=model_path,
        pytorch_model=module_scoped_subclassed_model,
        conda_env=None,
        pickle_module=pickle)

    mlflow_torch_pickle_load = mlflow_pytorch_pickle_module.load
    pickle_call_results = {
        "mlflow_torch_pickle_load_called": False,
    }

    def validate_mlflow_torch_pickle_load_called(*args, **kwargs):
        pickle_call_results["mlflow_torch_pickle_load_called"] = True
        return mlflow_torch_pickle_load(*args, **kwargs)

    log_messages = []

    def custom_warn(message_text, *args, **kwargs):
        log_messages.append(message_text % args % kwargs)

    with mock.patch("mlflow.pytorch.pickle_module.load") as mlflow_torch_pickle_load_mock,\
            mock.patch("mlflow.pytorch._logger.warning") as warn_mock:
        mlflow_torch_pickle_load_mock.side_effect = validate_mlflow_torch_pickle_load_called
        warn_mock.side_effect = custom_warn
        mlflow.pytorch.load_model(model_uri=model_path, pickle_module=mlflow_pytorch_pickle_module)

    assert all(pickle_call_results.values())
    assert any([
        "does not match the pickle module that was used to save the model" in log_message and
        pickle.__name__ in log_message and
        mlflow_pytorch_pickle_module.__name__ in log_message
        for log_message in log_messages
    ]) 

def __init__(self, name="some_curve"):
            """
            A CurveDB object has
            - name   = string to give the curve a name
            - pk     = integer to uniquely identify the curve (the database primary key)
            - data   = pandas.Series() object to hold any data
            - params = dict() with all kinds of parameters
            """
            self.logger = logging.getLogger(name=__name__)
            self.params = dict()
            x, y = np.array([], dtype=np.float), np.array([], dtype=np.float)
            self.data = (x, y)
            self.name = name 

def save(self):
            with open(os.path.join(self._dirname, str(self.pk) + self.file_extension),
                      'wb' if file_backend.__name__ == 'pickle' else 'w')\
                    as f:
                # wb is for compatibility with python 3
                # see http://stackoverflow.com/questions/5512811/builtins-typeerror-must-be-str-not-bytes
                data = [a.tolist() for a in self.data]
                file_backend.dump([self.pk, self.params, data], f, ) 

def read_msg(self, timeout=None):
        """
        Read data from com interface.

        :param timeout: timeout for reading data.
        :type timeout: float
        :return: (True, data) when reading is successful.
                 (False, None) when other side is closed.
                 (None, None) when reading is timeouted.
        """
        data = self._read_until_len(timeout)
        if data is None:
            return (None, None)
        if len(data) == 0:
            return (False, None)
        rdata = None
        try:
            cmd_len = int(data)
            rdata = ""
            rdata_len = 0
            while (rdata_len < cmd_len):
                rdata += self.stdin.read(cmd_len - rdata_len)
                rdata_len = len(rdata)
            rdataIO = StringIO(self.stdin.decode(rdata))
            unp = cPickle.Unpickler(rdataIO)
            if cPickle.__name__ == 'pickle':
                unp.find_class = _map_path
            else:
                unp.find_global = _map_path
            data = unp.load()
        except Exception as e:
            logging.error("ERROR data:%s rdata:%s" % (data, rdata))
            try:
                self.write_msg(remote_interface.MessengerError("Communication "
                                                               "failed.%s" % (e)))
            except OSError:
                pass
            self.flush_stdin()
            raise
        # Debugging commands.
        # if (isinstance(data, remote_interface.BaseCmd)):
        #    print data.func
        return (True, data)