Merge pull request #50 from pedohorse/environment-resolver-interface-refactor

Environment resolver interface refactor

Author: pedohorse

src/lifeblood/environment_resolver.py

@@ -18,21 +18,25 @@
 import inspect
 import pathlib
 import re
+import shutil
 from copy import deepcopy
 from semantic_version import Version, SimpleSpec
 from types import MappingProxyType
 from . import invocationjob, paths, logging
 from .config import get_config
 from .toml_coders import TomlFlatConfigEncoder
-from .process_utils import oh_no_its_windows
+from .process_utils import create_process, oh_no_its_windows
+from .exceptions import ProcessInitializationError
 
-from typing import Dict, Mapping, Optional, Type, Iterable
+from typing import Dict,  Iterable, List, Mapping, Optional, Type
 
 
 _resolvers: Dict[str, Type["BaseEnvironmentResolver"]] = {}  # this should be loaded from plugins
 
 
 def _populate_resolvers():
+    # TODO: This horrible logic needs refactoring. All plugins must be loaded with pluginloader
+    #  everything is tightly coupled here, interdependent... no planing
     for k, v in dict(globals()).items():
         if not inspect.isclass(v) \
                 or not issubclass(v, BaseEnvironmentResolver) \
@@ -53,7 +57,7 @@ class ResolutionImpossibleError(RuntimeError):
 
 class EnvironmentResolverArguments:
     """
-    this class objects specity requirements a task/invocation have for int's worker environment wrapper.
+    this class objects specify requirements a task/invocation have for it's worker environment wrapper.
     """
     def __init__(self, resolver_name=None, arguments: Optional[Mapping] = None):
         """
@@ -86,19 +90,24 @@ def remove_argument(self, name: str):
     def get_resolver(self):
         return get_resolver(self.__resolver_name)
 
-    def get_environment(self) -> "invocationjob.Environment":
-        return get_resolver(self.name()).get_environment(self.arguments())
+    async def get_environment(self) -> "invocationjob.Environment":
+        return await get_resolver(self.name()).get_environment(self.arguments())
 
     def serialize(self) -> bytes:
-        return json.dumps(self.__dict__).encode('utf-8')
+        return json.dumps({
+            '_EnvironmentResolverArguments__resolver_name': self.__resolver_name,
+            '_EnvironmentResolverArguments__args': self.__args,
+        }).encode('utf-8')
 
     async def serialize_async(self):
         return await asyncio.get_running_loop().run_in_executor(None, self.serialize)
 
     @classmethod
     def deserialize(cls, data: bytes):
         wrp = EnvironmentResolverArguments(None)
-        wrp.__dict__.update(json.loads(data.decode('utf-8')))
+        data_dict = json.loads(data.decode('utf-8'))
+        wrp.__resolver_name = data_dict['_EnvironmentResolverArguments__resolver_name']
+        wrp.__args = data_dict['_EnvironmentResolverArguments__args']
         return wrp
 
     @classmethod
@@ -107,7 +116,7 @@ async def deserialize_async(cls, data: bytes):
 
 
 class BaseEnvironmentResolver:
-    def get_environment(self, arguments: Mapping) -> "invocationjob.Environment":
+    async def get_environment(self, arguments: Mapping) -> "invocationjob.Environment":
         """
         this is the main reason for environment wrapper's existance.
         give it your specific arguments
@@ -118,17 +127,65 @@ def get_environment(self, arguments: Mapping) -> "invocationjob.Environment":
         """
         raise NotImplementedError()
 
+    async def create_process(self, arguments: Mapping, call_args: List[str], *, env: Optional[invocationjob.Environment] = None, cwd: Optional[str] = None) -> asyncio.subprocess.Process:
+        """
+        this should create process, maybe in a special way
+
+        :param arguments: EnvironmentResolverArguments for the resolver
+        :param call_args: what to call: process and arguments
+        :param env: optional environment to launch process in. If None - get_environment should be called
+        :param cwd: current working directory for the process
+        """
+        raise NotImplementedError()
+
+
+class BaseSimpleProcessSpawnEnvironmentResolver(BaseEnvironmentResolver):
+    async def create_process(self, arguments: Mapping, call_args: List[str], *, env: Optional[invocationjob.Environment] = None, cwd: Optional[str] = None) -> asyncio.subprocess.Process:
+        if env is None:
+            env = await self.get_environment(arguments)
+
+        if os.path.isabs(call_args[0]):
+            bin_path = call_args[0]
+        else:
+            bin_path = shutil.which(call_args[0], path=env.get('PATH', ''))
+        if bin_path is None:
+            raise ProcessInitializationError(f'"{call_args[0]}" was not found. Check environment resolver arguments and system setup')
+
+        if cwd is None:
+            cwd = os.path.dirname(bin_path)
 
-class TrivialEnvironmentResolver(BaseEnvironmentResolver):
+        return await create_process(call_args, env, cwd)
+
+
+class BaseSimpleProcessSpawnEnvironmentResolverWithPythonCheat(BaseSimpleProcessSpawnEnvironmentResolver):
+    # even though lifeblood worker runs with python interpreter - it does NOT expect
+    #  python to be available in run environment. User might want to have multiple versions of packaged python.
+    #  However, simple one user artist might not care, and would just want python code to work without any packages setup.
+    #  So to simplify the life of smaller setup users, this hack was introduced.
+    async def create_process(self, arguments: Mapping, call_args: List[str], *, env: Optional[invocationjob.Environment] = None, cwd: Optional[str] = None) -> asyncio.subprocess.Process:
+        """
+        This introduces path to sys.executable if no python is found in PATH, yet `python` is first arg in call_args
+        """
+        if env is None:
+            env = await self.get_environment(arguments)
+        if call_args[0] in ('python', 'python.exe') and shutil.which(call_args[0], path=env.get('PATH', '')) is None:
+            env.append('PATH', os.path.dirname(sys.executable))
+
+        return await super().create_process(arguments, call_args, env=env, cwd=cwd)
+
+
+class TrivialEnvironmentResolver(BaseSimpleProcessSpawnEnvironmentResolverWithPythonCheat):
     """
     trivial environment wrapper does nothing
     """
-    def get_environment(self, arguments: dict) -> "invocationjob.Environment":
+    async def get_environment(self, arguments: Mapping) -> "invocationjob.Environment":
         env = invocationjob.Environment(os.environ)
+        for key, value in arguments.items():
+            env[key] = value
         return env
 
 
-class StandardEnvironmentResolver(BaseEnvironmentResolver):
+class StandardEnvironmentResolver(BaseSimpleProcessSpawnEnvironmentResolverWithPythonCheat):
     """
     will initialize environment based on requested software versions and it's own config
     will raise ResolutionImpossibleError if he doesn't know how to resolve given configuration
@@ -162,7 +219,7 @@ def __init__(self):
             self.logger.error('environment resolver configs found, but all have errors! Aborting!')
             raise RuntimeError('all resolver configs are broken')
 
-    def get_environment(self, arguments: Mapping) -> "invocationjob.Environment":
+    async def get_environment(self, arguments: Mapping) -> "invocationjob.Environment":
         """
 
         :param arguments: are expected to be in format of package_name: version_specification
@@ -377,21 +434,21 @@ def main(args):
 
     opts = parser.parse_args(args)
 
+    logger = logging.get_logger('environment resolver')
+    logger.info('auto detecting packages...')
     packages = StandardEnvironmentResolver.autodetect_software()
     if opts.basepath:
         for basepath in opts.basepath.split(','):
             packages.update(StandardEnvironmentResolver.autodetect_software(basepath))
+    for pkgname, v in packages.items():
+        for verstr in v.keys():
+            logger.info(f'found {pkgname} : {verstr}')
 
     if opts.command == 'generate':
         config = get_config('standard_environment_resolver')
         if opts.output:
             config.override_config_save_location(opts.output)
         config.set_toml_encoder_generator(TomlFlatConfigEncoder)
-        logger = logging.get_logger('environment resolver')
-        logger.info('standard environment resolver is used, but no configuration found. auto generating configuration...')
-        for pkgname, v in packages.items():
-            for verstr in v.keys():
-                logger.info(f'found {pkgname} : {verstr}')
 
         if opts.override:  # do full config override
             config.set_option_noasync('packages', packages)
@@ -413,6 +470,7 @@ def main(args):
                 config.set_option_noasync('packages', conf_packages)
 
     elif opts.command == 'scan':
+        print('\n')
         for pkgname, stuff in packages.items():
             print(f'{pkgname}:')
             for ver, meta in stuff.items():

src/lifeblood/worker.py

@@ -14,7 +14,7 @@
 from .nethelpers import get_addr_to, get_localhost, get_hostname
 from .net_classes import WorkerResources
 from .worker_metadata import WorkerMetadata
-from .exceptions import ProcessInitializationError, WorkerNotAvailable, AlreadyRunning,\
+from .exceptions import WorkerNotAvailable, AlreadyRunning, \
     InvocationMessageWrongInvocationId, InvocationMessageAddresseeTimeout, InvocationCancelled
 from .worker_messsage_processor import WorkerMessageProcessor
 from .scheduler_message_processor import SchedulerWorkerControlClient
@@ -25,7 +25,7 @@
 from . import environment_resolver
 from .enums import WorkerType, WorkerState, ProcessPriorityAdjustment
 from .paths import config_path
-from .process_utils import create_process, kill_process_tree
+from .process_utils import kill_process_tree
 from .misc import event_set_context
 from .net_messages.address import AddressChain, DirectAddress
 from .net_messages.exceptions import MessageTransferError
@@ -344,10 +344,14 @@ async def run_task(self, task: InvocationJob, report_to: AddressChain):
             try:
                 if task.environment_resolver_arguments() is None:
                     config = get_config('worker')
-                    env = environment_resolver.get_resolver(config.get_option_noasync('default_env_wrapper.name', 'TrivialEnvironmentResolver'))\
-                        .get_environment(config.get_option_noasync('default_env_wrapper.arguments', {}))
+                    resolver = environment_resolver.get_resolver(config.get_option_noasync('default_env_wrapper.name', 'TrivialEnvironmentResolver'))
+                    resolver_arguments = config.get_option_noasync('default_env_wrapper.arguments', {})
                 else:
-                    env = task.environment_resolver_arguments().get_environment()
+                    env_res_args = task.environment_resolver_arguments()
+                    resolver = env_res_args.get_resolver()
+                    resolver_arguments = env_res_args.arguments()
+
+                env = await resolver.get_environment(resolver_arguments)
             except environment_resolver.ResolutionImpossibleError as e:
                 self.__logger.error(f'cannot run the task: Unable to resolve environment: {str(e)}')
                 raise
@@ -377,14 +381,11 @@ async def run_task(self, task: InvocationJob, report_to: AddressChain):
                 # TODO: proper child process priority adjustment should be done, for now it's implemented in constructor.
                 self.__running_process_start_time = time.time()
 
-                if os.path.isabs(args[0]):
-                    bin_path = args[0]
-                else:
-                    bin_path = shutil.which(args[0], path=env.get('PATH'))
-                if bin_path is None:
-                    raise ProcessInitializationError(f'"{args[0]}" was not found. Check environment resolver arguments and system setup')
-
-                self.__running_process: asyncio.subprocess.Process = await create_process(args, env, os.path.dirname(bin_path))
+                self.__running_process: asyncio.subprocess.Process = await resolver.create_process(
+                    resolver_arguments,
+                    args,
+                    env=env
+                )
             except Exception as e:
                 self.__logger.exception('task creation failed with error: %s' % (repr(e),))
                 raise

src/lifeblood_client/environment_resolver.py

@@ -8,7 +8,7 @@
 
 class EnvironmentResolverArguments:
     """
-    this class objects specity requirements a task/invocation have for int's worker environment wrapper.
+    this class objects specify requirements a task/invocation have for it's worker environment wrapper.
     """
     def __init__(self, resolver_name=None, arguments=None):
         """
@@ -30,10 +30,15 @@ def arguments(self):
         return self.__args
 
     def serialize(self):  # type: () -> bytes
-        return json.dumps(self.__dict__).encode('utf-8')
+        return json.dumps({
+            '_EnvironmentResolverArguments__resolver_name': self.__resolver_name,
+            '_EnvironmentResolverArguments__args': self.__args,
+        }).encode('utf-8')
 
     @classmethod
     def deserialize(cls, data):  # type: (bytes) -> EnvironmentResolverArguments
         wrp = EnvironmentResolverArguments(None)
-        wrp.__dict__.update(json.loads(data.decode('utf-8')))
+        data_dict = json.loads(data.decode('utf-8'))
+        wrp.__resolver_name = data_dict['_EnvironmentResolverArguments__resolver_name']
+        wrp.__args = data_dict['_EnvironmentResolverArguments__args']
         return wrp

src/lifeblood_testing_common/nodes_common.py

@@ -20,9 +20,9 @@
 from lifeblood.pluginloader import PluginNodeDataProvider
 from lifeblood.processingcontext import ProcessingContext
 from lifeblood.process_utils import oh_no_its_windows
-from lifeblood.environment_resolver import EnvironmentResolverArguments
+from lifeblood.environment_resolver import EnvironmentResolverArguments, BaseSimpleProcessSpawnEnvironmentResolver
 
-from typing import Any, Callable, Dict, List, Optional, Set, Union
+from typing import Any, Callable, Dict, List, Mapping, Optional, Set, Union
 
 
 plugin_data_provider = PluginNodeDataProvider()
@@ -34,17 +34,26 @@ def create_node(node_type: str, node_name: str, scheduler, node_id):
     return node
 
 
-class FakeEnvArgs(EnvironmentResolverArguments):
+class FakeResolver(BaseSimpleProcessSpawnEnvironmentResolver):
     def __init__(self, path_to_bin: str):
         super().__init__()
         self.__bin_path = Path(path_to_bin)
 
-    def get_environment(self):
+    async def get_environment(self, arguments: Mapping) -> "invocationjob.Environment":
         return Environment({**os.environ,
                             'PATH': os.pathsep.join((str(self.__bin_path), os.environ.get('PATH', ''))),
                             'PYTHONUNBUFFERED': '1'})
 
 
+class FakeEnvArgs(EnvironmentResolverArguments):
+    def __init__(self, path_to_bin: str):
+        super().__init__()
+        self.__bin_path = Path(path_to_bin)
+
+    def get_resolver(self):
+        return FakeResolver(self.__bin_path)
+
+
 class PseudoTaskPool:
     def children_ids_for(self, task_id, active_only=False) -> List[int]:
         raise NotImplementedError()

tests/test_envisonment_resolver.py

@@ -2,6 +2,7 @@
 import unittest
 from unittest import mock
 from lifeblood import environment_resolver
+from lifeblood_client import environment_resolver as client_environment_resolver
 from lifeblood.invocationjob import Environment
 from lifeblood.toml_coders import TomlFlatConfigEncoder
 import toml
@@ -21,7 +22,7 @@ def __set__(self, obj, val):
         self(obj, val)
 
 
-class StandardEnvResTest(unittest.TestCase):
+class StandardEnvResTest(unittest.IsolatedAsyncioTestCase):
     _stash = None
     @classmethod
     def setUpClass(cls) -> None:
@@ -33,10 +34,10 @@ def tearDownClass(cls) -> None:
         if cls._stash is not None:
             os.environ['LIFEBLOOD_CONFIG_LOCATION'] = cls._stash
 
-    def test_one(self):
+    async def test_one(self):
         ser = environment_resolver.StandardEnvironmentResolver()
         origenv = Environment(os.environ)
-        env = ser.get_environment({'package.houdini': '>18.0.0,<19.0.0'})
+        env = await ser.get_environment({'package.houdini': '>18.0.0,<19.0.0'})
 
         self.assertEqual(
             os.pathsep.join(["/path/to/hfs/bin", "/some/other/path/dunno", origenv.get('PATH', ''), "/whatever/you/want/to/append"]),
@@ -48,11 +49,11 @@ def test_one(self):
             env['PYTHONPATH']
         )
 
-    def test_two(self):
+    async def test_two(self):
         ser = environment_resolver.StandardEnvironmentResolver()
         origenv = Environment(os.environ)
-        env = ser.get_environment({'package.houdini': '>18.0.0,<19.0.0',
-                                   'package.assmouth': '~=2.3.2'})
+        env = await ser.get_environment({'package.houdini': '>18.0.0,<19.0.0',
+                                         'package.assmouth': '~=2.3.2'})
 
         self.assertEqual(
             os.pathsep.join(["/path/to/hfs/bin", "/some/other/path/dunno", "bananus", origenv.get('PATH', ''), "who/are/you", "/whatever/you/want/to/append"]),
@@ -149,3 +150,23 @@ def test_autodetect_win(self):
             print(result)
             self.assertIn('houdini.py3_10', result)
             self.assertEqual('C:\\Program Files\\Side Effects Software\\Houdini 19.5.640\\bin', result['houdini.py3_10']['19.5.640']['env']['PATH']['prepend'])
+
+
+class TestMainVsClientCompatibility(unittest.TestCase):
+    def test_serde1(self):
+        client_envarg = client_environment_resolver.EnvironmentResolverArguments('foobar', {'abc': 'qwe', 'def': 2.3, 'ghi': ['q', 2, 3.3, {'4': []}]})
+        data = client_envarg.serialize()
+        envarg = environment_resolver.EnvironmentResolverArguments.deserialize(data)
+
+        self.assertEqual(client_envarg.name(), envarg.name())
+        self.assertEqual(client_envarg.arguments(), envarg.arguments())
+        self.assertEqual(client_envarg.serialize(), envarg.serialize())
+
+    def test_serde1inv(self):
+        envarg = environment_resolver.EnvironmentResolverArguments('foobar', {'abc': 'qwe', 'def': 2.3, 'ghi': ['q', 2, 3.3, {'4': []}]})
+        data = envarg.serialize()
+        client_envarg = client_environment_resolver.EnvironmentResolverArguments.deserialize(data)
+
+        self.assertEqual(client_envarg.name(), envarg.name())
+        self.assertEqual(client_envarg.arguments(), envarg.arguments())
+        self.assertEqual(client_envarg.serialize(), envarg.serialize())