test cleanup to venvs

chrismallinson · chrismallinson · commit 14f9cc28f275 · 2025-10-18T13:58:36.000-04:00
diff --git a/.github/workflows/nightly-seed-grouping.yml b/.github/workflows/nightly-seed-grouping.yml
@@ -51,9 +51,9 @@ jobs:
             # rust-sdk
           ]
         include:
-        - concurrency-count-override: 8
-          sdk-name: python-sdk
-        
+          - concurrency-count-override: 8
+            sdk-name: python-sdk
+
     steps:
       # Note: this step is meant for when running on ubuntu-latest
       - name: Remove Bloat
@@ -63,7 +63,7 @@ jobs:
           # this might remove tools that are actually needed,
           # if set to "true" but frees about 6 GB
           tool-cache: false
-          
+
           # all of these default to true, but feel free to set to
           # "false" if necessary for your workflow
           android: true
@@ -99,14 +99,21 @@ jobs:
       - name: Save Seed Test Log to File
         id: save-seed-test-log-to-file
         run: |
-          # Default to 16 concurrent fixtures if not overridden. In the case of python-sdk, 
+          # Default to 16 concurrent fixtures if not overridden. In the case of python-sdk,
           # the virtual environments take up too much space so this needs to be limited
           pnpm seed:local test --generator ${{ matrix.sdk-name }} --parallel ${{ matrix.concurrency-count-override || 16 }} --allow-unexpected-failures > ${{ matrix.sdk-name }}-seed-test-log.txt
 
-      # CHRISM - consideration
-      # - name: Cleanup Virtualenvs
-      #     run: |
-      #       find . -name "fern-*" -type d | head -10 | xargs rm -rf
+      - name: Cleanup Virtualenvs
+        if: always() # Run even if tests fail
+        run: |
+          echo "Cleaning up Poetry virtualenvs to free disk space..."
+          # Remove Poetry virtualenv cache
+          rm -rf ~/.cache/pypoetry/virtualenvs/*
+          # Also clean any local .venv directories
+          find . -name ".venv" -type d -exec rm -rf {} + 2>/dev/null || true
+          # Clean poetry cache
+          poetry cache clear --all pypi 2>/dev/null || true
+          echo "Virtualenv cleanup completed"
 
       - name: Upload Seed Test Log as Artifact
         uses: actions/upload-artifact@v4
@@ -134,8 +141,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        sdk-name:
-          [
+        sdk-name: [
             # CHRISM - temp
             python-sdk
             # ruby-model,
diff --git a/generators/python/src/fern_python/cli/abstract_generator.py b/generators/python/src/fern_python/cli/abstract_generator.py
@@ -140,11 +140,13 @@ def generate_project(
             publisher.run_poetry_install()
             publisher.run_ruff_check_fix()
             publisher.run_ruff_format()
+            publisher.cleanup_virtualenvs()  # Clean up virtualenv after all operations
         elif output_mode_union.type == "publish":
             publisher.run_poetry_install()
             publisher.run_ruff_check_fix()
             publisher.run_ruff_format()
             publisher.publish_package(publish_config=output_mode_union)
+            publisher.cleanup_virtualenvs()  # Clean up virtualenv after all operations
 
         self.postrun(
             generator_exec_wrapper=generator_exec_wrapper,
@@ -197,6 +199,7 @@ def _publish(
             generator_config=generator_config,
         )
         publisher.publish_package(publish_config=publish_config)
+        publisher.cleanup_virtualenvs()  # Clean up virtualenv after publishing
 
     def _write_files_for_github_repo(
         self, project: Project, output_mode: GithubOutputMode, write_unit_tests: bool
diff --git a/generators/python/src/fern_python/cli/publisher.py b/generators/python/src/fern_python/cli/publisher.py
@@ -51,6 +51,44 @@ def run_poetry_install(self) -> None:
             safe_command="poetry install",
         )
 
+    def cleanup_virtualenvs(self) -> None:
+        """Remove Poetry virtualenvs to free up disk space"""
+        try:
+            # Get virtualenv info to identify which one was created
+            env_info_result = subprocess.run(
+                ["poetry", "env", "info", "--path"],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                cwd=self._generator_config.output.path,
+                check=False,  # Don't fail if no venv exists
+                text=True
+            )
+
+            if env_info_result.returncode == 0 and env_info_result.stdout.strip():
+                venv_path = env_info_result.stdout.strip()
+                print(f"Removing virtualenv at: {venv_path}")
+
+                # Remove the specific virtualenv
+                self._run_command(
+                    command=["poetry", "env", "remove", "--all"],
+                    safe_command="poetry env remove --all",
+                )
+                print("Virtualenv cleanup completed successfully")
+            else:
+                print("No virtualenv found to clean up")
+
+        except Exception as e:
+            # Don't fail the entire generation process if cleanup fails
+            print(f"Warning: Virtualenv cleanup failed: {e}")
+            self._generator_exec_wrapper.send_update(
+                logging.GeneratorUpdate.factory.log(
+                    logging.LogUpdate(
+                        level=logging.LogLevel.DEBUG,
+                        message=f"Virtualenv cleanup warning: {e}"
+                    )
+                )
+            )
+
     def publish_package(
         self,
         *,
diff --git a/packages/seed/src/commands/test/script-runner/DockerScriptRunner.ts b/packages/seed/src/commands/test/script-runner/DockerScriptRunner.ts
@@ -51,6 +51,47 @@ export class DockerScriptRunner extends ScriptRunner {
         }
     }
 
+    public async cleanup({ taskContext, id }: { taskContext: TaskContext; id: string }): Promise<void> {
+        if (this.skipScripts) {
+            return;
+        }
+
+        const workDir = id.replace(":", "_");
+
+        for (const script of this.scripts) {
+            taskContext.logger.debug(`Cleaning up fixture ${id} in container ${script.containerId}`);
+
+            // Clean up the fixture directory and any virtualenvs
+            const cleanupCommands = [
+                `rm -rf /${workDir}`,
+                // Clean up Poetry virtualenvs created during this fixture's execution
+                `find . -name '.venv*' -type d -exec rm -rf {} + 2>/dev/null || true`,
+                // Clean up any virtualenvs in the Poetry cache for this specific project
+                `rm -rf ./*-py* 2>/dev/null || true`,
+                // Clean up Poetry cache
+                `poetry cache clear --all pypi 2>/dev/null || true`
+            ];
+
+            const cleanupScript = cleanupCommands.join(" && ");
+
+            const cleanupCommand = await loggingExeca(
+                taskContext.logger,
+                "docker",
+                ["exec", script.containerId, "/bin/sh", "-c", cleanupScript],
+                {
+                    doNotPipeOutput: false,
+                    reject: false // Don't fail if cleanup has issues
+                }
+            );
+
+            if (cleanupCommand.failed) {
+                taskContext.logger.warn(`Cleanup warning for fixture ${id}: ${cleanupCommand.stderr}`);
+            } else {
+                taskContext.logger.debug(`Successfully cleaned up fixture ${id}`);
+            }
+        }
+    }
+
     protected async initialize(): Promise<void> {
         await this.startContainers(this.context);
     }
diff --git a/packages/seed/src/commands/test/script-runner/LocalScriptRunner.ts b/packages/seed/src/commands/test/script-runner/LocalScriptRunner.ts
@@ -40,6 +40,27 @@ export class LocalScriptRunner extends ScriptRunner {
         // No containers to stop for local execution
     }
 
+    public async cleanup({ taskContext, id }: { taskContext: TaskContext; id: string }): Promise<void> {
+        if (this.skipScripts) {
+            return;
+        }
+
+        taskContext.logger.debug(`Cleaning up fixture ${id} (local execution)`);
+
+        // For local execution, clean up Poetry virtualenvs in the output directory
+        // This is less critical since local runs use temporary directories that are cleaned up anyway
+        // But we can still clean up Poetry caches to free space
+        try {
+            await loggingExeca(taskContext.logger, "poetry", ["cache", "clear", "--all", "pypi"], {
+                doNotPipeOutput: false,
+                reject: false
+            });
+            taskContext.logger.debug(`Successfully cleaned up Poetry cache for fixture ${id}`);
+        } catch (error) {
+            taskContext.logger.warn(`Cleanup warning for fixture ${id}: ${error}`);
+        }
+    }
+
     protected async initialize(): Promise<void> {
         // No initialization needed for local execution
     }
diff --git a/packages/seed/src/commands/test/script-runner/ScriptRunner.ts b/packages/seed/src/commands/test/script-runner/ScriptRunner.ts
@@ -37,6 +37,7 @@ export abstract class ScriptRunner {
 
     public abstract run({ taskContext, id, outputDir }: ScriptRunner.RunArgs): Promise<ScriptRunner.RunResponse>;
     public abstract stop(): Promise<void>;
+    public abstract cleanup({ taskContext, id }: { taskContext: TaskContext; id: string }): Promise<void>;
 
     protected abstract initialize(): Promise<void>;
 }
diff --git a/packages/seed/src/commands/test/test-runner/TestRunner.ts b/packages/seed/src/commands/test/test-runner/TestRunner.ts
@@ -263,6 +263,14 @@ export abstract class TestRunner {
             scriptStopwatch.stop();
             metrics.compileTime = scriptStopwatch.duration();
 
+            // Clean up virtualenvs after script execution (success or failure)
+            try {
+                await this.scriptRunner?.cleanup({ taskContext, id });
+            } catch (cleanupError) {
+                taskContext.logger.warn(`Cleanup failed for fixture ${id}: ${cleanupError}`);
+                // Don't fail the test due to cleanup issues
+            }
+
             if (scriptResponse?.type === "failure") {
                 return {
                     type: "failure",

Original file line number	Diff line number	Diff line change
`@@ -37,6 +37,7 @@ export abstract class ScriptRunner {`
`37`	`37`
`38`	`38`	`public abstract run({ taskContext, id, outputDir }: ScriptRunner.RunArgs): Promise<ScriptRunner.RunResponse>;`
`39`	`39`	`public abstract stop(): Promise<void>;`
	`40`	`+ public abstract cleanup({ taskContext, id }: { taskContext: TaskContext; id: string }): Promise<void>;`
`40`	`41`
`41`	`42`	`protected abstract initialize(): Promise<void>;`
`42`	`43`	`}`