Complete localFarm implementation on tests

Author: Alxiice

.codecov.yml

@@ -50,3 +50,4 @@ ignore:
   - "docs/"
   - "scripts/"
   - "bin/"
+  - "localFarm/"  # For now ignore localFarm as it has no coverage yet

bin/meshroom_compute

@@ -72,6 +72,7 @@ graph.update()
 if args.node:
     # Execute the node
     node = graph.findNode(args.node)
+    node.updateStatusFromCache()
     submittedStatuses = [Status.RUNNING]
     if not args.extern:
         # If running as "extern", the task is supposed to have the status SUBMITTED.
@@ -80,7 +81,7 @@ if args.node:
 
     if not node._chunksCreated:
         print(f"Error: Node {node} has been submitted before chunks have been created." \
-               "See file: \"{node.nodeStatusFile}\".")
+              f"See file: \"{node.nodeStatusFile}\".")
         sys.exit(-1)
 
     if node._isInputNode():

bin/meshroom_createChunks

@@ -137,7 +137,7 @@ if chunksToProcess:
     node.postprocess()
     node.restoreLogger()
 else:
-    logging.info(f"[MeshroomCreateChunks] -> create job to process chunks {node.chunks}")
+    logging.info(f"[MeshroomCreateChunks] -> create job to process chunks {[c for c in node.chunks]}")
     submitter.createChunkTask(node, graphFile=args.graphFile, cache=args.cache, 
                               forceStatus=args.forceStatus, forceCompute=args.forceCompute)
 

localfarm/__init__.py

@@ -4,14 +4,16 @@
 Local Farm : A simple local job runner
 """
 
+from __future__ import annotations  # For forward references in type hints
+
 import logging
 import json
 import socket
 import logging
 import uuid
-from collections import defaultdict, deque
+from collections import defaultdict
 from pathlib import Path
-from typing import Dict, List
+from typing import Dict, List, Generator
 
 logging.basicConfig(
     level=logging.INFO,
@@ -30,6 +32,7 @@ def __init__(self, root):
 
     def connect(self):
         """Connect to the backend"""
+        print("Connect to farm located at", self.root)
         if self.tcpPortFile.exists():
             try:
                 port = int(self.tcpPortFile.read_text())
@@ -68,7 +71,7 @@ def _call(self, method, **params):
         finally:
             sock.close()
 
-    def submit_job(self, job):
+    def submit_job(self, job: Job):
         """ Submit the job to the farm """
         # Create the job
         createdJob = self._call("create_job", name=job.name)
@@ -83,22 +86,23 @@ def submit_job(self, job):
                     raise RuntimeError(f"Parent task {parentTask.name} not created yet")
                 deps.append(tasksCreated[parentTask])
             createdTask = self._call("create_task", 
-                jid=jid, name=task.name, command=task.command, metadata=task.metadata, dependencies=deps)
+                jid=jid, name=task.name, command=task.command, 
+                metadata=task.metadata, dependencies=deps, env=task.env)
             tasksCreated[task] = createdTask["tid"]
         # Submit the job
         self._call("submit_job", jid=jid)
         return {"jid": jid}
 
     def create_additional_task(self, jid, tid, task):
         """ Create new task in an existing job """
-        metadata = {}
         createdTask = self._call("expand_task", 
-            jid=jid, name=task.name, command=task.command, metadata=task.metadata, parentTid=tid)
+            jid=jid, name=task.name, command=task.command, 
+            metadata=task.metadata, parentTid=tid, env=task.env)
         return {"tid": createdTask["tid"]}
 
     def get_job_infos(self, jid):
         """Get job status"""
-        return self._call("get_job_status", jid=jid)
+        return self._call("get_job_infos", jid=jid)["result"]
 
     def pause_job(self, jid):
         """Pause a job"""
@@ -132,9 +136,15 @@ def restart_task(self, jid, tid):
         """Restart a task"""
         return self._call("restart_task", jid=jid, tid=tid)
 
-    def list_jobs(self):
+    def list_jobs(self) -> list:
         """List all jobs"""
-        return self._call("list_jobs")
+        return self._call("list_jobs")["jobs"]
+
+    def getJobStatus(self, jid: int) -> dict:
+        for job in self.list_jobs():
+            if job["jid"] == jid:
+                return job
+        return {}
 
     def ping(self):
         """Check if backend is alive"""
@@ -146,12 +156,13 @@ def ping(self):
 
 
 class Task:
-    def __init__(self, name, command, metadata):
+    def __init__(self, name, command, metadata=None, env=None):
         self.uid = str(uuid.uuid1())
         self.name = name
         self.command = command
-        self.metadata = metadata
-    
+        self.metadata = metadata or {}
+        self.env = env or {}
+
     def __repr__(self):
         return f"<Task {self.name}|{self.uid}>"
 
@@ -165,13 +176,17 @@ def __init__(self, name):
         self.tasks: Dict[str, Task] = {}
         self.dependencies: Dict[str: List[str]] = defaultdict(set)
         self.reverseDependencies: Dict[str: List[str]] = defaultdict(set)
+        self._engine: LocalFarmEngine = None
+    
+    def setEngine(self, engine: LocalFarmEngine):
+        self._engine = engine
 
     def addTask(self, task):
         if task.name in self.tasks:
             raise ValueError(f"Task {task} already exists in job")
         self.tasks[task.uid] = task
 
-    def addTaskDependency(self, task, dependsOn):
+    def addTaskDependency(self, task: Task, dependsOn: Task):
         if task.uid not in self.tasks:
             raise ValueError(f"Task {task} not found in job")
         if dependsOn.uid not in self.tasks:
@@ -217,7 +232,7 @@ def exploreTask(taskUid, taskParents=None):
                 return True
         return False
 
-    def tasksDFS(self) -> List[Task]:
+    def tasksDFS(self) -> Generator[Task]:
         """
         Return tasks in topological order (dependencies before dependents).
         Tasks closer to roots appear first.
@@ -243,10 +258,13 @@ def exploreTask(task: str, currentLevel=0):
             for task in tasks:
                 yield task
 
-    def submit(self, farmPath):
-        farm = LocalFarmEngine(farmPath)
-        result = farm.submit_job(self)
-        return result
+    def submit(self, engine: LocalFarmEngine = None):
+        engine = engine or self._engine
+        if engine:
+            result = engine.submit_job(self)
+            return result
+        else:
+            raise ValueError("No LocalFarmEngine set for this job")
 
 
 def test():

meshroom/submitters/localFarm/localFarm.py renamed to localfarm/localFarm.py

@@ -23,7 +23,7 @@
 import threading
 from socketserver import BaseRequestHandler, ThreadingTCPServer
 
-FARM_MAX_PARALLEL_TASKS = 2
+FARM_MAX_PARALLEL_TASKS = 10
 MAX_BYTES_REQUEST = 4096  # 8192 / 65536 if needed
 
 PathLike = Union[str, Path]
@@ -48,14 +48,15 @@ class Status(Enum):
 
 
 class Task:
-    def __init__(self, jid: str, tid: str, label: str, command: str, metadata: dict, jobDir: PathLike):
+    def __init__(self, jid: str, tid: str, label: str, command: str, metadata: dict, jobDir: PathLike, env: dict = None):
         self.jid: str = jid
         self.tid: str = tid
         self.parentTids = []  # Tasks that must be completed before this one
         self.childTids = []   # Task that depend on this one
         self.label: str = label
         self.command: str = command
-        self.metadata: dict = metadata
+        self.metadata: dict = metadata or {}
+        self.env: dict = env or {}
         self.taskDir: Path = Path(jobDir) / "tasks"
         self.taskDir.mkdir(parents=True, exist_ok=True)
         self.status: Status = Status.NONE
@@ -73,6 +74,7 @@ def to_dict(self):
             "label": self.label,
             "command": self.command,
             "metadata": self.metadata,
+            "env": self.env,
             "status": self.status.name,
             "created_at": self.created_at.isoformat(),
             "started_at": self.started_at.isoformat() if self.started_at else None,
@@ -300,14 +302,7 @@ def processJobs(self):
                     # Check if process finished
                     returncode = task.process.poll()
                     if returncode is not None:
-                        task.finished_at = datetime.now()
-                        task.return_code = returncode
-                        if returncode == 0:
-                            task.status = Status.SUCCESS
-                            logger.info(f"Task {task.tid} completed")
-                        else:
-                            task.status = Status.ERROR
-                            logger.error(f"Task {task.tid} failed with code {returncode}")
+                        self.finishTask(task, returncode)
 
             # Check if job is complete
             if any(t.status in [Status.ERROR, Status.STOPPED, Status.KILLED] for t in job.tasks):
@@ -342,20 +337,48 @@ def startTask(self, task: Task):
         task.status = Status.RUNNING
         task.started_at = datetime.now()
         # Create log file
+        additional_env = {
+            "LOCALFARM_CURRENT_JID": str(task.jid),
+            "LOCALFARM_CURRENT_TID": str(task.tid),
+            "MR_LOCAL_FARM_PATH": str(self.root)
+        }
+        additional_env.update(task.env)
+        process_env = os.environ.copy()
+        process_env.update(additional_env)
         try:
+
             with open(task.logFile, "w") as log:
+                log.write(f"# ========== Starting task {task.tid} at {task.started_at.isoformat()}" \
+                          f" (command=\"{task.command}\") ==========\n")
+                log.write(f"# Additional env variables:\n")
+                for _k, _v in additional_env.items():
+                    log.write(f"# - {str(_k)}={str(_v)}\n")
+                log.write(f"\n")
                 task.process = subprocess.Popen(
                     task.command,
                     # shlex.split(task.command),
                     stdout=log,
                     stderr=log,
                     cwd=task.taskDir,
+                    env=process_env,
                     shell=True
                 )
         except Exception as e:
             logger.error(f"Failed to start task {task.tid}: {e}")
             task.status = "error"
             task.finished_at = datetime.now()
+    
+    def finishTask(self, task: Task, returncode: int):
+        task.finished_at = datetime.now()
+        task.return_code = returncode
+        if returncode == 0:
+            task.status = Status.SUCCESS
+            logger.info(f"Task {task.tid} completed")
+        else:
+            task.status = Status.ERROR
+            logger.error(f"Task {task.tid} failed with code {returncode}")
+        with open(task.logFile, "a") as log:
+            log.write(f"\n# ========== Task {task.tid} finished at {task.finished_at.isoformat()} ==========\n")
 
     def cleanup(self):
         logger.info("Cleaning up...")
@@ -393,15 +416,15 @@ def create_job(self, name):
             logger.info(f"Created job {jid}")
             return {"success": True, "jid": jid}
 
-    def create_task(self, jid, name, command, metadata, dependencies):
+    def create_task(self, jid, name, command, metadata, dependencies, env=None):
         """Add a task to a job"""
         with self.lock:
             if jid not in self.jobs:
                 return {"success": False, "error": "Job not found"}
             job = self.jobs[jid]
             job.lastJid += 1
             tid = job.lastJid
-            task = Task(jid, tid, name, command, metadata, job.jobDir)
+            task = Task(jid, tid, name, command, metadata, job.jobDir, env=env)
             job.tasks.append(task)
             for parentTid in dependencies:
                 parentTask = next((t for t in job.tasks if t.tid == parentTid), None)
@@ -412,14 +435,15 @@ def create_task(self, jid, name, command, metadata, dependencies):
             logger.info(f"Added task {tid} to job {jid}")
             return {"success": True, "tid": tid}
 
-    def expand_task(self, jid, name, command, metadata, parentTid):
+    def expand_task(self, jid, name, command, metadata, parentTid, env=None):
         with self.lock:
             if jid not in self.jobs:
+                logger.info(f"Available jobs: {list(self.jobs.keys())}")
                 return {"success": False, "error": "Job not found"}
             job = self.jobs[jid]
             job.lastJid += 1
             tid = job.lastJid
-            task = Task(jid, tid, name, command, metadata, job.jobDir)
+            task = Task(jid, tid, name, command, metadata, job.jobDir, env=env)
             task.status = Status.SUBMITTED
             job.tasks.append(task)
             parentTask = next((t for t in job.tasks if t.tid == parentTid), None)
@@ -456,7 +480,6 @@ def get_job_infos(self, jid):
             if jid not in self.jobs:
                 return {'success': False, "error": "Job not found"}
             job = self.jobs[jid]
-            logger.info(f"Job infos : {job.to_dict()}")
             return {"success": True, "result": job.to_dict()}
 
     def pause_job(self, jid):
@@ -538,7 +561,7 @@ def list_jobs(self):
         with self.lock:
             return {
                 "success": True,
-                "jobs": {jid: job.to_dict() for jid, job in self.jobs.items()}
+                "jobs": [job.to_dict() for job in self.jobs.values()]
             }
 
 
@@ -609,6 +632,9 @@ def main(root):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description='Execute a Graph of processes.')
-    parser.add_argument('farm_root', type=str, help='Root path for the farm.')
+    parser.add_argument('--root', type=str, required=False, help='Root path for the farm.')
     args = parser.parse_args()
-    main(args.farm_root)
+    root = args.root
+    if not root:
+        root = os.getenv("MR_LOCAL_FARM_PATH", os.path.join(os.path.expanduser("~"), ".local_farm"))
+    main(root)