fix (#79)

command launch fixed

Author: abhinavsriva

pyproject.toml

@@ -1,6 +1,10 @@
+[build-system]
+requires = ["setuptools>=61", "wheel"]
+build-backend = "setuptools.build_meta"
+
 [project]
 name = "traceml-ai"
-version = "0.2.9"
+version = "0.2.10"
 
 description = "TraceML: Lightweight training runtime health monitor."
 authors = [{ name = "Abhinav Srivastav", email = "abhinav@traceopt.ai" }]

src/traceml/__init__.py

@@ -6,7 +6,7 @@
     trace_step,
 )
 
-__version__ = "0.2.9"
+__version__ = "0.2.10"
 
 __all__ = [
     "__version__",

src/traceml/cli.py

@@ -14,6 +14,7 @@
 
 import msgspec
 
+from traceml.runtime.launch_context import LaunchContext
 from traceml.runtime.session import get_session_id
 from traceml.utils.ast_analysis import analyze_script, build_code_manifest
 
@@ -200,6 +201,7 @@ def write_run_manifest(
     nproc_per_node: int,
     history_enabled: bool,
     status: str,
+    launch_cwd: str,
     aggregator_dir: Optional[Path] = None,
     db_path: Optional[Path] = None,
     extra: Optional[Dict[str, Any]] = None,
@@ -233,6 +235,7 @@ def write_run_manifest(
             "tcp_port": int(tcp_port),
             "nproc_per_node": int(nproc_per_node),
             "history_enabled": bool(history_enabled),
+            "launch_cwd": str(Path(launch_cwd).resolve()),
         },
         "paths": {
             "session_root": str(session_root),
@@ -389,10 +392,14 @@ def _handler(signum: int, _frame: Any) -> None:
     signal.signal(signal.SIGTERM, _handler)
 
 
-def start_aggregator_process(env: Dict[str, str]) -> subprocess.Popen:
+def start_aggregator_process(
+    env: Dict[str, str], cwd: str
+) -> subprocess.Popen:
     """Start the TraceML aggregator as a separate process.
 
-    Stdout/stderr are inherited so Rich output and tracebacks remain visible.
+    The subprocess cwd is set explicitly so all child processes inherit a
+    deterministic working directory rather than depending on ambient shell
+    state.
     """
     aggregator_path = (
         Path(__file__).parent / "aggregator" / "aggregator_main.py"
@@ -404,18 +411,29 @@ def start_aggregator_process(env: Dict[str, str]) -> subprocess.Popen:
 
     cmd = [sys.executable, str(aggregator_path)]
     print("[TraceML] Launching TraceML aggregator:", " ".join(cmd))
-    return subprocess.Popen(cmd, env=env, start_new_session=True)
+    return subprocess.Popen(
+        cmd,
+        env=env,
+        cwd=cwd,
+        start_new_session=True,
+    )
 
 
 def start_training_process(
-    train_cmd: list[str], env: Dict[str, str]
+    train_cmd: list[str], env: Dict[str, str], cwd: str
 ) -> subprocess.Popen:
     """Start the training process in a new process group.
 
-    Stdout/stderr are inherited so user logs and tracebacks remain visible.
+    The subprocess cwd is set explicitly so worker processes see the same
+    working directory the user launched TraceML from.
     """
     print("[TraceML] Launching TraceML executor:", " ".join(train_cmd))
-    return subprocess.Popen(train_cmd, env=env, start_new_session=True)
+    return subprocess.Popen(
+        train_cmd,
+        env=env,
+        cwd=cwd,
+        start_new_session=True,
+    )
 
 
 def launch_process(script_path: str, args: argparse.Namespace) -> None:
@@ -430,6 +448,7 @@ def launch_process(script_path: str, args: argparse.Namespace) -> None:
     5. Keep training as the primary process; aggregator may fail open
     6. On shutdown, terminate child process groups and update the manifest
     """
+
     env = os.environ.copy()
     env["PYTHONUNBUFFERED"] = "1"
 
@@ -452,6 +471,10 @@ def launch_process(script_path: str, args: argparse.Namespace) -> None:
     env["TRACEML_NPROC_PER_NODE"] = str(args.nproc_per_node)
     env["TRACEML_HISTORY_ENABLED"] = "0" if args.no_history else "1"
 
+    launch_context = LaunchContext.capture()
+    env.update(launch_context.to_env())
+    execution_cwd = launch_context.launch_cwd
+
     session_id = env["TRACEML_SESSION_ID"]
     session_root = Path(args.logs_dir).resolve() / session_id
     aggregator_dir = session_root / "aggregator"
@@ -474,6 +497,7 @@ def launch_process(script_path: str, args: argparse.Namespace) -> None:
         nproc_per_node=args.nproc_per_node,
         history_enabled=not args.no_history,
         status="starting",
+        launch_cwd=execution_cwd,
         aggregator_dir=aggregator_dir,
         db_path=db_path,
         extra=(
@@ -495,7 +519,11 @@ def launch_process(script_path: str, args: argparse.Namespace) -> None:
             str(script_path),
             *script_args,
         ]
-        train_proc = start_training_process(train_cmd=train_cmd, env=env)
+        train_proc = start_training_process(
+            train_cmd=train_cmd,
+            env=env,
+            cwd=execution_cwd,
+        )
         install_shutdown_handlers(
             lambda: (train_proc, None), manifest_path=manifest_path
         )
@@ -530,7 +558,7 @@ def launch_process(script_path: str, args: argparse.Namespace) -> None:
         f"(ui={args.mode}, profile={env['TRACEML_PROFILE']})"
     )
     try:
-        agg_proc = start_aggregator_process(env=env)
+        agg_proc = start_aggregator_process(env=env, cwd=execution_cwd)
     except FileNotFoundError as exc:
         print(f"[TraceML] ERROR: {exc}", file=sys.stderr)
         update_run_manifest(manifest_path, status="failed")
@@ -558,7 +586,11 @@ def launch_process(script_path: str, args: argparse.Namespace) -> None:
     print("[TraceML] Aggregator ready.")
     update_run_manifest(manifest_path, status="running")
 
-    train_proc = start_training_process(train_cmd=train_cmd, env=env)
+    train_proc = start_training_process(
+        train_cmd=train_cmd,
+        env=env,
+        cwd=execution_cwd,
+    )
 
     while True:
         train_rc = train_proc.poll()

src/traceml/runtime/executor.py

@@ -30,6 +30,10 @@
 from pathlib import Path
 from typing import Any, Dict, Optional, Union
 
+from traceml.runtime.launch_context import (
+    LaunchContext,
+    script_execution_context,
+)
 from traceml.runtime.runtime import TraceMLRuntime
 from traceml.runtime.settings import TraceMLSettings, TraceMLTCPSettings
 from traceml.utils.shared_utils import EXECUTION_LAYER
@@ -344,21 +348,21 @@ def stop_runtime(
 
 def run_user_script(script_path: str, script_args: list[str]) -> None:
     """
-    Execute the user script in-process using runpy.
+    Execute the user script in-process using ``runpy``.
 
-    This intentionally does not spawn another subprocess so that:
-    - hooks attach to the real Python objects
-    - stack traces remain meaningful
-    - execution context matches the user's script as closely as possible
-
-    Global sys.argv is restored afterward to avoid leaking state.
+    TraceML preserves the user's original launch cwd while also exposing the
+    script directory as ``sys.path[0]`` so local imports and relative file
+    access behave like a normal ``python`` or ``torchrun`` script launch.
     """
-    old_argv = sys.argv[:]
-    try:
-        sys.argv = [script_path, *script_args]
-        runpy.run_path(script_path, run_name="__main__")
-    finally:
-        sys.argv = old_argv
+    launch_context = LaunchContext.from_env()
+    resolved_script_path = str(Path(script_path).resolve())
+
+    with script_execution_context(
+        script_path=resolved_script_path,
+        script_args=script_args,
+        launch_context=launch_context,
+    ):
+        runpy.run_path(resolved_script_path, run_name="__main__")
 
 
 def report_crash(cfg: Dict[str, Any], error: BaseException) -> None:

src/traceml/runtime/launch_context.py

@@ -0,0 +1,104 @@
+"""
+Helpers for preserving the user's original launch context.
+
+TraceML should instrument user code without changing how that code resolves:
+
+- current working directory
+- sibling imports
+- script argv
+
+This module centralizes that behavior so the CLI launcher and runtime
+executor stay small and consistent.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, Iterator
+
+TRACEML_LAUNCH_CWD_ENV = "TRACEML_LAUNCH_CWD"
+
+
+@dataclass(frozen=True)
+class LaunchContext:
+    """Execution context captured from the user's original CLI invocation.
+
+    Attributes
+    ----------
+    launch_cwd:
+        The directory where the user invoked ``traceml``.
+    """
+
+    launch_cwd: str
+
+    @classmethod
+    def capture(cls) -> "LaunchContext":
+        """Capture the current user-facing launch context.
+
+        This should be called in the main CLI process before spawning any
+        subprocesses so child processes can faithfully reproduce the user's
+        original execution environment.
+        """
+        return cls(launch_cwd=str(Path.cwd().resolve()))
+
+    @classmethod
+    def from_env(cls) -> "LaunchContext":
+        """Load a previously captured launch context from environment vars.
+
+        Falls back to the current process cwd when no explicit launch cwd has
+        been provided. This keeps behavior safe for tests and direct executor
+        invocation.
+        """
+        raw = os.environ.get(TRACEML_LAUNCH_CWD_ENV, "").strip()
+        if raw:
+            return cls(launch_cwd=str(Path(raw).resolve()))
+        return cls.capture()
+
+    def to_env(self) -> Dict[str, str]:
+        """Serialize this launch context into environment variables."""
+        return {TRACEML_LAUNCH_CWD_ENV: self.launch_cwd}
+
+
+@contextmanager
+def script_execution_context(
+    *,
+    script_path: str,
+    script_args: list[str],
+    launch_context: LaunchContext,
+) -> Iterator[None]:
+    """Temporarily apply Python script execution semantics.
+
+    Behavior
+    --------
+    - ``sys.argv`` looks like a direct script launch
+    - ``sys.path[0]`` points at the script directory
+    - process cwd matches the user's original launch cwd
+
+    This combination most closely matches normal ``python train.py`` or
+    ``torchrun train.py`` behavior from the user's project directory.
+    """
+    resolved_script_path = str(Path(script_path).resolve())
+    script_dir = str(Path(resolved_script_path).parent)
+
+    old_argv = sys.argv[:]
+    old_path = sys.path[:]
+    old_cwd = os.getcwd()
+
+    try:
+        sys.argv = [resolved_script_path, *script_args]
+
+        if sys.path:
+            sys.path[0] = script_dir
+        else:
+            sys.path = [script_dir]
+
+        os.chdir(launch_context.launch_cwd)
+        yield
+    finally:
+        os.chdir(old_cwd)
+        sys.argv = old_argv
+        sys.path = old_path